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: