024号文書

主にプログラミング

ABC105 D - Candy Distribution

ABC

問題

https://atcoder.jp/contests/abc105/tasks/abc105_d

解法

解のパターンを以下の通り分割します。

  • パターン1: 1番目の箱のみ使うパターン
  • パターン2: 2番目の箱、および、必要ならば1番目の箱を使うパターン
  • パターン3: 3番目の箱、および、必要ならば1, 2番目の箱を使うパターン
  • ...
  • パターンN: N番目の箱、および、必要ならば1, 2, ..., N-1番目の箱を使うパターン

上記のパターン分けは直和分割なので、最終的に求める解は上記の各パターンの解の和です。 パターンiの解を x_iで表します

x_1 は簡単に求められます。 A_1 がMの倍数ならば、1, そうでなければ0です。

次に、 x_2 を求めます。

  • A_2 がMの倍数か
  • A_1 + A_2 がMの倍数か

をチェックし、結果に応じて、0, 1, 2のいずれかの値になります。

x_3 を求めます。 同様に、3通りの計算を実施し、Mの倍数かチェックしてもよいのですが、そろそろ一般化も考えてみましょう。 この方法だと、 x_i はi通りの計算が必要となり、計算量が大きすぎます。

ちょっと唐突ですが、 A_1 + A_2 + A_3 の先頭から何個かの項を除去し、Mの倍数にできないかを考えます。 A_1 + A_2 + A_3 をMで割った余りをRとします。 0, A_1, A_1 + A_2 のうち、余りがRであるものの個数が x_3 です。 余りがRであるものを数え上げるのが大変そうですが、先頭からk要素の和の余りをハッシュで記憶しておけば、一般の場合においても、ほぼ定数オーダーの計算で x_i を計算できます。 （ハッシュの初期値は {0: 1} となる。ハッシュの各要素(k, v)は余りがkとなるものがv個あることを意味する）

最後に、上記のアルゴリズムの動作の具体例を書きます。 サンプル2で遊んでみます。

  • 29%17=12はハッシュに含まれません。パターン1の解は0です。ハッシュに12を追加します。
    • 更新後のハッシュ: {0: 1, 12: 1}
    • ここまでの合計値: 12
  • (12+7=19)%17=2はハッシュに含まれません。パターン2の解は0です。ハッシュに2を追加します。
    • 更新後のハッシュ: {0: 1, 2: 1, 12: 1}
    • ここまでの合計値: 2
  • 2+5=7はハッシュに含まれません。パターン3の解は0です。ハッシュに7を追加します。
    • 更新後のハッシュ: {0: 1, 2: 1, 7: 1, 12: 1}
    • ここまでの合計値: 7
  • 7+7=14はハッシュに含まれません。パターン4の解は0です。ハッシュに14を追加します。
    • 更新後のハッシュ: {0: 1, 5: 1, 7: 1, 12: 1, 14: 1}
    • ここまでの合計値: 14
  • (14+9=23)%17=6はハッシュに含まれません。パターン5の解は0です。ハッシュに6を追加します。
    • 更新後のハッシュ: {0: 1, 5: 1, 6: 1, 7: 1, 12: 1, 14: 1}
    • ここまでの合計値: 6
  • (6+51=57)%17=6はハッシュに含まれます。パターン6の解は1です。ハッシュに6を追加します（既に存在するので値を1から2に増やします）。
    • 更新後のハッシュ: {0: 1, 5: 1, 6: 2, 7: 1, 12: 1, 14: 1}
    • ここまでの合計値: 6
  • 6+7=13はハッシュに含まれません。パターン7の解は0です。ハッシュに13を追加します。
    • 更新後のハッシュ: {0: 1, 5: 1, 6: 2, 7: 1, 12: 1, 13: 0, 14: 1}
    • ここまでの合計値: 13
  • (13+13=26)%17=9はハッシュに含まれません。パターン8の解は0です。ハッシュに9を追加します。
    • 更新後のハッシュ: {0: 1, 5: 1, 6: 2, 7: 1, 9: 1, 12: 1, 13: 0, 14: 1}
    • ここまでの合計値: 9
  • (9+8=17)%17=0はハッシュに含まれます。パターン9の解は1です。ハッシュに0を追加します。
    • 更新後のハッシュ: {0: 2, 5: 1, 6: 2, 7: 1, 9: 1, 12: 1, 13: 0, 14: 1}
    • ここまでの合計値: 0
  • (0+55=55)%17=4はハッシュに含まれません。パターン10の解は0です。ハッシュに4を追加します。
    • 更新後のハッシュ: {0: 1, 4: 1, 5: 1, 6: 2, 7: 1, 9: 1, 12: 1, 13: 0, 14: 1}
    • ここまでの合計値: 4
  • (4+42=46)%17=12はハッシュに含まれます。パターン11の解は1です。ハッシュに12を追加します。
    • 更新後のハッシュ: {0: 1, 4: 1, 5: 1, 6: 2, 7: 1, 9: 1, 12: 2, 13: 0, 14: 1}
    • ここまでの合計値: 12
  • (12+9=21)%17=4はハッシュに含まれます。パターン12の解は1です。ハッシュに4を追加します。
    • 更新後のハッシュ: {0: 1, 4: 2, 5: 1, 6: 2, 7: 1, 9: 1, 12: 2, 13: 0, 14: 1}
    • ここまでの合計値: 4
  • (4+81=85)%17=0はハッシュに含まれます。パターン13の解は2です。ハッシュに0を追加します。
    • 更新後のハッシュ: {0: 2, 4: 2, 5: 1, 6: 2, 7: 1, 9: 1, 12: 2, 13: 0, 14: 1}
    • ここまでの合計値: 0

以上より、各パターンの解を合計すると6を得ます。

ACしたコード

from collections import defaultdict
N, M = map(int, input().split())
A = list(map(int, input().split()))
acc = 0
s = defaultdict(int)
s[0] = 1
ans = 0
for a in A:
    acc = (acc + a) % M
    ans += s[acc]
    s[acc] += 1
print(ans)