公式

E - Subsegment Sum Finder 解説 by physics0523


原案: physics0523

以下の問題が解ければよいです。なぜなら、以下の問題で \(X=L\) とした場合の答えから \(X=R+1\) とした場合の答えを引けば本問題の答えになるからです。

連続部分列であって、総和が \(X\) 以上のものを求める。

そして、累積和を使うと変換後の問題は以下のように言い換えられます。

\(S_0=0, S_k = A_1+A_2+\dots+A_k (k \ge 1)\) とする。
各 \(i=1,2,\dots,N\) について、 \(0 \le j < i\) かつ \(S_i-S_j \ge X\) であるものの数を数える。

これは、 \(i\) の昇順に処理すると、これまでに出た \(S_j \le S_i-X\) なる \(j\) の個数を足し合わせることで求められます。
例えば、最初に全ての \(S_i\) を求めて座標圧縮したうえで、 \(S_i\) の昇順に既に出たかどうかを管理する fenwick tree の区間和などを利用すればこれを実現できます。

本解法の時間計算量は \(O(N \log N)\) です。

実装例 (C++):

#include<bits/stdc++.h>
#include<atcoder/fenwicktree>

using namespace std;
using namespace atcoder;

using ll=long long;
using pl=pair<ll,ll>;

ll f(ll x, vector<ll> &s, vector<pl> &vp){
  ll n=vp.size()-1;
  vector<ll> inv(n+1);
  for(ll i=0;i<=n;i++){
    inv[vp[i].second]=i;
  }

  fenwick_tree<ll> fw(n+1);
  ll res=0;
  for(ll i=0;i<=n;i++){
    ll pos=(lower_bound(vp.begin(),vp.end(),pl{s[i]-x+1,-1})-vp.begin());
    res+=fw.sum(0,pos);
    fw.add(inv[i],1);
  }
  return res;
}

int main(){
  ll n,l,r;
  cin >> n >> l >> r;
  vector<ll> s(n+1,0);
  vector<pl> vp={{0,0}};
  for(ll i=1;i<=n;i++){
    ll a;
    cin >> a;
    s[i]=s[i-1]+a;
    vp.push_back({s[i],i});
  }
  sort(vp.begin(),vp.end());
  cout << f(l,s,vp)-f(r+1,s,vp) << "\n";
  return 0;
}