概要

この問題は、与えられた \(N\) 個の整数からなる数列において、その連続する部分列の和がちょうど \(K\) となる箇所がいくつあるかを数え上げる問題です。

考察

素朴なアプローチ

まず、すべての連続部分列 \((l, r)\) の組み合わせを調べる方法を考えます。 左端 \(l\) を \(1\) から \(N\) まで、右端 \(r\) を \(l\) から \(N\) まで動かしてその合計を計算すると、組み合わせの数は \(O(N^2)\) 通りとなります。今回の制約では \(N \leq 2 \times 10^5\) であるため、計算回数が最大で \(4 \times 10^{10}\) 回程度になり、実行時間制限（通常 2 秒）に間に合いません。

累積和の活用

区間の和を効率的に計算するために、累積和を利用します。 数列の先頭から \(i\) 番目までの要素の合計を \(S_i\) とすると（ただし \(S_0 = 0\)）、区間 \([l, r]\) の和は以下のように表せます。 $\(\sum_{i=l}^{r} A_i = S_r - S_{l-1}\)$

今回の目標は \(S_r - S_{l-1} = K\) となる \((l-1, r)\) のペアの数を探すことです。 この式を変形すると、以下のようになります。 $\(S_{l-1} = S_r - K\)$

これは、「現在の累積和 \(S_r\) に到達したとき、過去に累積和が \(S_r - K\) であった箇所が何回あったか」を数えればよいことを意味します。

アルゴリズム

ハッシュマップ（Pythonでは dict または defaultdict）を使用して、これまでに現れた累積和の頻度を記録しながら数列を 1 回走査します。

1. 累積和の頻度を記録する辞書 prefix_sums を用意し、prefix_sums[0] = 1 と初期化します（これは「何も選んでいない状態」の累積和が 0 であることを示します）。
2. 現在の累積和 current_sum を 0 とし、答え ans を 0 とします。
3. 数列の各要素 \(x\) について、以下の操作を繰り返します。
   • current_sum に \(x\) を加える。
   • target = current_sum - K を計算する。
   • prefix_sums[target] の値を ans に加算する（過去に target という累積和が何回出たかを確認）。
   • prefix_sums[current_sum] の値を 1 増やす。
4. 最終的な ans を出力します。

計算量

• 時間計算量: \(O(N)\)
  • 数列を 1 回走査し、各ステップでのハッシュマップへのアクセス（挿入・検索）は平均 \(O(1)\) で行えるため、全体で \(O(N)\) となります。
• 空間計算量: \(O(N)\)
  • 最悪の場合、すべての累積和が異なる値になり、ハッシュマップに \(N+1\) 個の要素が格納されるため、 \(O(N)\) となります。

実装のポイント

• prefix_sums[0] = 1 の意味: もし累積和がちょうど \(K\) になった場合（つまり \(S_r = K\)）、\(S_r - K = 0\) となります。このとき、過去の累積和 \(S_0 = 0\) をカウントできるように、あらかじめ 0 の出現回数を 1 としておく必要があります。

• 大きな値の扱い: \(K\) や \(A_i\) の値が負であったり、非常に大きな値（\(10^9\) 単位）であったりしますが、累積和とハッシュマップを用いる手法であれば、これらを特別に意識することなく一貫して処理できます。

  ソースコード

import sys
from collections import defaultdict

def solve():
    # 標準入力からすべてのデータを読み込み、スペースで分割する
    data = sys.stdin.read().split()
    if not data:
        return
    
    # N: 数列の要素数, K: 目標とする合計値
    N = int(data[0])
    K = int(data[1])
    
    # 数列 A の各要素を整数に変換するイテレータ
    A = map(int, data[2:])
    
    # 累積和の出現回数を記録する辞書
    # prefix_sums[s] は累積和が s となる箇所の数
    # 累積和が 0 である状態（何も選んでいない状態）をあらかじめ 1 回として数えておく
    prefix_sums = defaultdict(int)
    prefix_sums[0] = 1
    
    current_sum = 0
    ans = 0
    
    # 各要素を順番に走査し、現在の累積和を更新しながら条件を満たす区間を数える
    for x in A:
        current_sum += x
        
        # 現在の累積和を S_j とすると、S_j - S_i = K となるような S_i が
        # 過去に何回出現したかを確認する。これは S_i = S_j - K と同値。
        target = current_sum - K
        ans += prefix_sums[target]
        
        # 現在の累積和を記録する
        prefix_sums[current_sum] += 1
        
    # 結果を出力
    print(ans)

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-preview によって生成されました。