公式

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;
}

投稿日時:
最終更新: