けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

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

AtCoder AtCoder400点 ABC-D 茶色diff NoviSteps2Q Zero-Sum Ranges setやmap:検索や削除が速いこと 計算量改善:setやmapの活用 連想配列(setやmap) 数列 累積和 解空間:O(N^2)個の区間 解空間:O(N^2)通りの選択肢 負値が本質的に効く問題 数え上げ問題

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

問題概要

数列 A_{1}, A_{2}, \dots, A_{N} と整数 K が与えられる。

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

制約

  • 1 \le N \le 2 \times 10^{5}
  • -10^{9} \le A_{i} \le 10^{9}
  • -10^{15} \le K \le 10^{15}

考えたこと

0-indexed で考える。

数列 A_{0}, A_{1}, \dots, A_{N-1} の累積和を S_{0}, S_{1}, \dots, S_{N} としよう。このとき、

A_{l} + A_{l+1} + \dots + A_{r-1} = S_{r} - S_{l}

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

N+1 個の数からなる数列 S_{0}, S_{1}, \dots, S_{N} が与えられる。

  • 0 \le l \lt r \le N
  • S_{r} - S_{l} = K

を満たすような (l, r) の組の個数を求めよ。

こうなれば、解きやすくなる。 r = 0, 1, 2, \dots, N に対して、

  • S_{l} = S_{r} - K かつ 0 \le l \lt r

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

  • nums[x]：値 x の個数

計算量は O(N \log N) となる。

コード

#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;
}