Official
D - Count Interval Editorial by sugarrr
区間和に関する問題では、累積和を取ることが典型テクニックです。
\(0\leq i \leq N\) に対して \(S_i = \sum_{j=1}^{i}A_j\) と累積和をとると、問題は以下のようになります
以下の条件を満たす整数の組 \((l,r)\) はいくつあるか?
- \(1\leq l\leq r\leq N\)
- \(S_r - S_{l-1}=K\)
これに基づいて一旦 \(O(N^2)\) 解を考えると、次のようになります
ans = 0
for r in 1 .. n
for l in 1 .. r
if S[l-1] == S[r] - K
ans += 1
出力(ans)
擬似コード \(3\) 行目から \(5\) 行目の \(l\) のループを削除して計算量を減らすことを考えます。
そのためには、\(1\leq l\leq r\) かつ \(S[l-1] = S[r]-K\) を満たす \(l\) の数が高速に求められれば良いです。
これは、C++での std::map
のような連想配列を用いるなどして求めることができます。
以上でこの問題を \(O(N \log N)\) で解くことができました。
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: