概要

数列から連続する部分列を選び、その合計がちょうど \(K\) になる選び方の数を求める問題です。累積和とハッシュマップを組み合わせることで効率的に解けます。

考察

素朴なアプローチとその問題点

最も単純な方法は、すべての連続部分列を試すことです。 - 始点 \(l\) を \(1\) から \(N\) まで - 終点 \(r\) を \(l\) から \(N\) まで - 各区間の合計を計算

この方法では、区間の選び方が \(O(N^2)\) 通りあり、各区間の合計計算に \(O(N)\) かかると全体で \(O(N^3)\) となります。\(N = 2 \times 10^5\) では間に合いません。

累積和による高速化

累積和 \(S\) を導入します： - \(S[0] = 0\) - \(S[i] = A[1] + A[2] + \cdots + A[i]\)

すると、区間 \([l, r]\) の合計は \(S[r] - S[l-1]\) で \(O(1)\) で求まります。

重要な気づき

区間 \([l, r]\) の合計が \(K\) になる条件は： $\(S[r] - S[l-1] = K\)$

これを変形すると： $\(S[l-1] = S[r] - K\)$

つまり、各 \(r\) に対して、\(S[r] - K\) と等しい累積和がいくつ存在するか を高速に数えられれば良いのです。

具体例

数列 \(A = [1, 2, 3, -2, 2]\)、\(K = 3\) の場合： - \(S = [0, 1, 3, 6, 4, 6]\)

\(r = 2\)（\(S[2] = 3\)）のとき、\(S[r] - K = 0\) となる \(S[l-1]\) を探すと、\(S[0] = 0\) が該当。 これは区間 \([1, 2]\) の和が \(3\) であることを意味します。

アルゴリズム

1. ハッシュマップ prefix_count を用意し、prefix_count[0] = 1 で初期化（\(S[0] = 0\) に対応）
2. 累積和 current_sum を \(0\) で初期化
3. 配列 \(A\) を先頭から順に走査：
   • current_sum += A[i]（これが \(S[i+1]\) に相当）
   • target = current_sum - K を計算
   • prefix_count[target] を答えに加算（target と等しい過去の累積和の個数）
   • prefix_count[current_sum] を \(1\) 増やす
4. 答えを出力

例: A = [1, 2, -1, 2], K = 3

i=0: current_sum=1, target=-2, count+=0, prefix_count={0:1, 1:1}
i=1: current_sum=3, target=0,  count+=1, prefix_count={0:1, 1:1, 3:1}
i=2: current_sum=2, target=-1, count+=0, prefix_count={0:1, 1:1, 3:1, 2:1}
i=3: current_sum=4, target=1,  count+=1, prefix_count={0:1, 1:1, 3:1, 2:1, 4:1}

答え: 2（区間[1,2]と区間[2,4]）

計算量

• 時間計算量: \(O(N)\)
  • 配列を一度走査し、各要素に対してハッシュマップの操作（平均 \(O(1)\)）を行う
• 空間計算量: \(O(N)\)
  • ハッシュマップに最大 \(N+1\) 個の累積和を保存

実装のポイント

• 初期化: prefix_count[0] = 1 を忘れずに設定する。これは「何も選ばない」状態（\(S[0] = 0\)）に対応し、先頭から始まる部分列を正しくカウントするために必要

• 順序: 現在の累積和を prefix_count に追加する前に、target の出現回数を数える。これにより、自分自身を重複カウントしない

• 負の数の扱い: 要素に負の数が含まれるため、尺取り法は使えない。ハッシュマップによる方法が適切

  ソースコード

from collections import defaultdict

def solve():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    
    # 累積和を使用
    # S[i] = A[0] + A[1] + ... + A[i-1] (S[0] = 0)
    # 区間 [l, r] の和 = S[r+1] - S[l]
    # これが K になる条件: S[r+1] - S[l] = K
    # つまり S[l] = S[r+1] - K
    
    # 各累積和の値が何回出現したかを記録
    prefix_count = defaultdict(int)
    prefix_count[0] = 1  # S[0] = 0
    
    current_sum = 0
    count = 0
    
    for i in range(N):
        current_sum += A[i]  # これが S[i+1] に相当
        # S[i+1] - K となる累積和がいくつあるか
        target = current_sum - K
        count += prefix_count[target]
        # 現在の累積和を記録
        prefix_count[current_sum] += 1
    
    print(count)

solve()

---

この解説は claude4.5opus によって生成されました。