概要

連続する部分列の総重量が \(K\) の倍数となる選び方の個数を、累積和と剰余の性質を利用して効率的に数え上げる問題です。

考察

素朴なアプローチ

すべての連続部分列を試す方法を考えると、始点 \(l\) と終点 \(r\) の組み合わせは \(O(N^2)\) 通りあります。さらに各部分列の和を計算すると、全体で \(O(N^3)\) または累積和を使っても \(O(N^2)\) かかり、\(N \leq 2 \times 10^5\) では間に合いません。

重要な気づき：累積和と剰余の関係

累積和 \(S[i] = A_1 + A_2 + \cdots + A_i\)（\(S[0] = 0\)）を定義すると、連続部分列 \([l+1, r]\) の和は \(S[r] - S[l]\) と表せます。

この和が \(K\) の倍数になる条件は： $\(S[r] - S[l] \equiv 0 \pmod{K}\)$

これを変形すると： $\(S[r] \equiv S[l] \pmod{K}\)$

つまり、累積和を \(K\) で割った余りが同じ2つのインデックスを選べば、その間の部分列の和は \(K\) の倍数になるということです。

具体例

\(N = 4, K = 3, A = [1, 2, 3, 6]\) の場合：

\(i\)	\(S[i]\)	\(S[i] \mod 3\)
0	0	0
1	1	1
2	3	0
3	6	0
4	12	0

余りが 0 のインデックスは \(\{0, 2, 3, 4\}\) の4個あるので、この中から2つ選ぶ組み合わせ \(_4C_2 = 6\) 通りが答えに寄与します。

アルゴリズム

1. 累積和の \(\mod K\) を順番に計算していく
2. 各余りの値が何回出現したかを辞書（ハッシュマップ）で管理
3. 新しい累積和の余り \(r\) を計算したとき、すでに同じ余り \(r\) が出現していた回数を答えに加算
4. その後、余り \(r\) の出現回数を1増やす

なぜこれで正しく数えられるか？

累積和 \(S[i]\) を順に見ていき、\(S[i] \mod K = r\) となったとき、それ以前に同じ余り \(r\) を持つ累積和が \(c\) 個あれば、\(S[i]\) とそれら \(c\) 個それぞれとペアを作れます。各ペアが1つの有効な部分列に対応します。

初期状態で count[0] = 1 としているのは、\(S[0] = 0\)（何も選ばない状態）を表しています。

計算量

• 時間計算量: \(O(N)\)
  • 配列を1回走査し、各要素に対して辞書への \(O(1)\) のアクセスを行う
• 空間計算量: \(O(\min(N, K))\)
  • 辞書に格納される余りの種類は最大で \(\min(N+1, K)\) 種類

実装のポイント

• defaultdict(int) を使うことで、まだ出現していない余りに対しても自動的に 0 が返される

• \(K\) が最大 \(10^9\) と大きいため、配列ではなく辞書を使う必要がある（配列だとメモリが足りない）

• 累積和を \(K\) で割った余りだけを保持すればよく、累積和そのものを配列で持つ必要はない

• count[0] = 1 の初期化を忘れると、先頭から始まる部分列を数え漏らすので注意

  ソースコード

from collections import defaultdict

def solve():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    
    # 累積和の mod K を計算
    # S[i] = (A[0] + A[1] + ... + A[i-1]) mod K
    # 連続部分列 [l, r) の和が K の倍数 ⇔ S[r] ≡ S[l] (mod K)
    
    # 累積和の mod K が同じものの個数を数える
    count = defaultdict(int)
    count[0] = 1  # 空の累積和（インデックス0の前）
    
    prefix_sum_mod = 0
    result = 0
    
    for i in range(N):
        prefix_sum_mod = (prefix_sum_mod + A[i]) % K
        # 同じ mod K を持つ以前の累積和の個数を加算
        result += count[prefix_sum_mod]
        count[prefix_sum_mod] += 1
    
    print(result)

solve()

---

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