Official

D - Count Interval Editorial by en_translator


One of the typical approaches for problems on intervals is to compute the cumulative sums.
Defining the cumulative sums \(S_i = \sum_{j=1}^{i}A_j\) for \(0\leq i \leq N\), the problem can be interpreted as follows:

How many pairs of integers \((l,r)\) are there that satisfy the conditions below?

  • \(1\leq l\leq r\leq N\)
  • \(S_r - S_{l-1}=K\)

This leads to a temporary \(O(N^2)\) solution:

ans = 0
for r in 1 .. n
    for l in 1 .. r
        if S[l-1] == S[r] - K
            ans += 1
output(ans)

How can we reduce the complexity by eliminating the loop from the \(3\)-th through \(5\)-th lines?

What we want to obtain fast is the number of \(l\) such that \(1\leq l\leq r\) and \(S[l-1] = S[r]-K\). This can be achieved with an associative array, like std::map in C++.

Therefore, the problem has been solved in a total of \(O(N \log N)\) time.

Sample code in C++

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i,l,r) for(ll i=l;i<=r;i++)

signed main(){
    ll n,k;cin>>n>>k;
    vector<ll>a(n);rep(i,0,n-1)cin>>a[i];
    vector<ll>s(n+1);rep(i,0,n-1)s[i+1]=s[i]+a[i];
    map<ll,ll>mp;
    ll ans=0;
    rep(r,1,n){
        mp[s[r-1]]++;
        ans += mp[s[r]-k];
    }
    cout<<ans<<endl;
    return 0;
}

posted:
last update: