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\) の倍数になる区間(先頭からの区間)も自然に処理できます。
アルゴリズム
- 累積和の剰余を順番に計算していきます。
- 各ステップで、現在の剰余が過去に何度出てきたかを調べ、その分だけ答えに加算します。
- 現在の剰余の出現回数をインクリメントします。
- 最初に剰余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: