AtCoder ABC 233 D - Count Interval (2Q, 茶色, 400 点)

「区間の値の和」を見たら、累積和をとろう！！

• 問題へのリンク

問題概要

数列 と整数 が与えられる。

数列の連続する区間であって、その総和が に一致するものが何個あるかを求めよ。

制約

• 
• 
• 

考えたこと

0-indexed で考える。

数列 の累積和を としよう。このとき、

となる。よって、もとの問題は次のようになる。

---

個の数からなる数列 が与えられる。

• 
• 

を満たすような の組の個数を求めよ。

---

こうなれば、解きやすくなる。 に対して、

• かつ

を満たすような の個数を合算していけばよい。これをやるためには、次の連想配列 (C++ ならば map) を動的に管理していけばよい。

• nums[x]：値 の個数

計算量は となる。

コード

#include <bits/stdc++.h>
using namespace std;

int main() {
    long long N, K;
    cin >> N >> K;
    vector<long long> A(N), S(N+1, 0);
    map<long long, long long> nums;
    for (int i = 0; i < N; i++) {
        cin >> A[i];
        S[i+1] = S[i] + A[i];
    }

    long long res = 0;
    for (int i = 0; i <= N; i++) {
        if (nums.count(S[i] - K)) res += nums[S[i] - K];
        nums[S[i]]++;
    }
    cout << res << endl;
}