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: