概要

\(N\) 個の要素からなる数列の中で、連続する部分列の和が \(K\) の倍数となる箇所の総数を求める問題です。

考察

素朴なアプローチ

連続する部分列の選び方は、左端を \(i\)、右端を \(j\) とすると全部で \(O(N^2)\) 通りあります。それぞれについて和を計算して \(K\) で割り切れるか判定すると、全体で \(O(N^3)\) または累積和を使っても \(O(N^2)\) の計算量がかかります。今回の制約は \(N \leq 2 \times 10^5\) であるため、これでは制限時間に間に合いません。

累積和と余りの性質

連続部分列の和を効率的に扱うために、累積和を利用します。 \(S_i\) を最初から \(i\) 番目までのお菓子の重さの総和とすると、 \(i+1\) 番目から \(j\) 番目までのお菓子の総和は \(S_j - S_i\) と表せます。

この和が \(K\) の倍数であるという条件は、以下の式で表せます： $\((S_j - S_i) \equiv 0 \pmod K\)\( これを変形すると： \)\(S_j \equiv S_i \pmod K\)\( つまり、**「累積和を \)K\( で割った余りが等しいペア \)(i, j)\( を選ぶ」**問題に言い換えることができます（ただし \)0 \leq i < j \leq N$）。

効率的な数え上げ

累積和の余りが同じになるインデックスが \(c\) 個あるとき、その中から 2 つを選ぶ組み合わせの数は \({}_c\mathrm{C}_2 = \frac{c(c-1)}{2}\) 通りです。 全ての余り（\(0\) から \(K-1\)）についてこの値を合計すれば、答えを \(O(N)\) で求めることができます。

アルゴリズム

1. 累積和の計算: 配列を先頭から走査し、現在の累積和を \(K\) で割った余りを計算していきます。
2. 余りの出現回数を記録: ハッシュマップ（Pythonでは collections.Counter や辞書）を使い、各余りが何回出現したかをカウントします。
   • 注意: 累積和が 0（最初のお菓子も選んでいない状態）を初期値としてカウントに含める必要があります。
3. 組み合わせの計算: 各余りの出現回数を \(c\) としたとき、\(\frac{c(c-1)}{2}\) を計算し、その総和を求めます。

計算量

• 時間計算量: \(O(N)\)
  • 数列を 1 回走査するのに \(O(N)\)、ハッシュマップの操作に平均 \(O(1)\) かかるため、全体で \(O(N)\) となります。
• 空間計算量: \(O(N)\)
  • 最悪の場合、累積和の余りがすべて異なるため、ハッシュマップに最大 \(N+1\) 個の要素を格納する分の空間が必要です。

実装のポイント

• \(K\) の大きさ: \(K\) は最大 \(10^9\) と非常に大きいため、配列でカウントしようとするとメモリ不足になります。必ず連想配列（辞書）を使用してください。

• 初期値の扱い: 何も選んでいない状態の累積和 \(S_0 = 0\) をあらかじめカウント（remainder_counts[0] = 1）しておくのを忘れないようにしましょう。これを忘れると、先頭から始まる部分列（例：\(A_1 + A_2\) が \(K\) の倍数の場合など）が正しく数えられません。

  ソースコード

import sys
from collections import Counter

def solve():
    # 入力を取得
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    n = int(input_data[0])
    k = int(input_data[1])
    a = list(map(int, input_data[2:]))
    
    # 累積和の余りをカウントする
    # S_0 = 0 なので、初期値として 0 を 1 つカウントしておく
    remainder_counts = Counter()
    remainder_counts[0] = 1
    
    current_sum = 0
    for x in a:
        current_sum = (current_sum + x) % k
        remainder_counts[current_sum] += 1
    
    # 同じ余りを持つ累積和のペアの数を計算する
    # 余りが同じ S_i, S_j (i < j) があれば、その間の区間の和は K の倍数
    ans = 0
    for count in remainder_counts.values():
        if count >= 2:
            ans += count * (count - 1) // 2
            
    print(ans)

if __name__ == "__main__":
    solve()

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