Official

D - お菓子の分配 / Distribution of Sweets Editorial by admin

Qwen3-Coder-480B(非推奨)

概要

連続するお菓子の区間で、その総重量が \(K\) の倍数になるものの個数を求めます。

考察

この問題では、連続した部分列の和が \(K\) の倍数になる組み合わせの数を数える必要があります。

素朴な方法として、全ての区間 \([l, r]\) を全探索し、それぞれの和を計算して \(K\) の倍数かどうか判定する方法がありますが、これは \(O(N^2)\) かかり、制約 \(N \leq 2 \times 10^5\) に対しては明らかに時間内に終わりません(TLE)。

重要な観察は、累積和を使うことです。
累積和 \(S_i = A_1 + A_2 + \cdots + A_i\) を用いると、区間 \([l, r]\) の和は \(S_r - S_{l-1}\) と表せます。
この和が \(K\) の倍数である条件は、

\[ (S_r - S_{l-1}) \equiv 0 \pmod{K} \]

つまり、

\[ S_r \equiv S_{l-1} \pmod{K} \]

となります。
つまり、「累積和の剰余が等しい2点を選ぶ」ことで、その間の区間和が \(K\) の倍数になることが分かります。

したがって、累積和を前から計算しながら、これまでに同じ剰余値が何回出現したかを記録しておけば、その時点で終わる有効な区間の数が求められます。

また、最初に剰余0のカウントを1にしておくことで、累積和そのものが \(K\) の倍数になる区間(先頭からの区間)も自然に処理できます。

アルゴリズム

  1. 累積和の剰余を順番に計算していきます。
  2. 各ステップで、現在の剰余が過去に何度出てきたかを調べ、その分だけ答えに加算します。
  3. 現在の剰余の出現回数をインクリメントします。
  4. 最初に剰余0の出現回数を1として初期化しておきます(空の区間に対応)。

例えば入力が以下の場合:

N=5, K=3
A = [1, 2, 3, 4, 5]

累積和の剰余は順に: - \(1 \bmod 3 = 1\) - \((1+2) \bmod 3 = 0\) - \((1+2+3) \bmod 3 = 0\) - \((1+2+3+4) \bmod 3 = 1\) - \((1+2+3+4+5) \bmod 3 = 0\)

剰余ごとに出現回数を管理しながら組み合わせを数えることで、計算量を削減できます。

計算量

  • 時間計算量: \(O(N)\)
  • 空間計算量: \(O(K)\) (剰余の種類数は最大 \(K\) だが、実際には出現する剰余の種類数に依存)

※実装では defaultdict を使っているので、出現する剰余の種類数に比例したメモリしか使いません。

実装のポイント

  • 累積和の剰余を逐次計算する際に、剰余0の初期値を1にしておくこと(最初の要素から始まる区間を考慮)。
  • 区間は「1個以上」なので、空の区間は含めないが、計算上は便宜的に扱う必要がある。
  • Pythonの sys.stdin.read を使うことで高速な入力処理を行う(特に \(N\) が大きい場合に重要)。
## ソースコード

```python
import sys
from collections import defaultdict

def main():
    input = sys.stdin.read
    data = input().split()
    
    N = int(data[0])
    K = int(data[1])
    A = list(map(int, data[2:]))

    # 累積和の剰余を計算
    cumsum_mod = 0
    mod_count = defaultdict(int)
    mod_count[0] = 1  # 空の累積和に対応するもの
    
    count = 0
    for i in range(N):
        cumsum_mod = (cumsum_mod + A[i]) % K
        # 同じ剰余を持つ以前の累積和の数だけ、Kの倍数になる区間が存在する
        count += mod_count[cumsum_mod]
        mod_count[cumsum_mod] += 1
        
    print(count)

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。

posted:
last update: