E - Subsegment Sum Finder 解説 by en_translator
Original proposer: physics0523
It suffices to solve the following problem, because subtracting the answer to the following problem for \(X=R+1\) from that for \(X=L\) yields the answer to the original problem.
Find the number of subarrays whose sum is \(X\) or greater.
Using prefix sums, the problem can be rephrased as follows:
Let \(S_0=0, S_k = A_1+A_2+\dots+A_k (k \ge 1)\).
For each \(i=1,2,\dots,N\), find the number of indices \(j\) such that \(0 \le j < i\) and \(S_i-S_j \ge X\).
By scanning in ascending order of \(i\), this can be done by summing up the number of indices \(j\) such that \(S_j \le S_i-X\) inspected so far.
This can be done by, for example, precalculating all \(S_i\), applying coordinate compression, and using a Fenwick tree to manage whether \(S_i\) has appeared already in ascending order of \(S_i\).
The time complexity of this solution is \(O(N \log N)\).
Sample code (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;
}