AtCoder Beginner Contest 164 D 数学よく分からないマンが解説を加えてみる

もはやAtCoder Mathematical Contestだという声の上がっている表題の問題について、

以下かつっぱさんの解説が華麗すぎてすぐには分からなかったので、数学がよく分からない人でもなんとなくわかった気になれるよう、解説を加えてみます。

www.youtube.com

考え方

文字列Sが"60574038"だったとします。
このとき、1文字目から4文字目の"6057"が割り切れるかを考えてみます。

Aを右端から"6057"の左端までの値、Bを右端から"6057"の右端手前までの値と考えます。つまり

A = 60574038
B = 4038

です。
このとき、AからBを取り除いたA-B(6057)が2019で割り切れるかどうかは

(A-B) / (10**4) % 2019 = 0

と表現できますが、ここで10**nと2019は互いに素なので、上の式は

(A-B) % 2019 = 0

と等価です。これは

A % 2019 = B % 2019

と等価です。つまり、2019で割った時のAとBの余りが同じなら、A-Bは2019で割り切れます。

このような(A,B)の組み合わせの数を計算すれば良いので、
余りが同じAやBがn個あったとき、その組み合わせのnC2、つまりn*(n-1)/2を計算すれば良いことになります。

従って本問題は、
まず余りの値ごとに、Sの中で余りが同じになる値の数を数え（余りが0の値が何個、余りが1の値が何個、、、）、
次に余りの値ごとに、組み合わせの数を計算し、合計した値が答え、
となります。

コード

コードは以下の通り。かつっぱさんのコード丸パクリですが、コメントで説明を加えます。

s = input()[::-1] # 右からn桁の数についてループを回していきたいので、文字列を反転させる

sum_of_digits = 0 # 余りの値
cnts = [0] * 2019 # 余りの値ごとに余りが同じ値の数を格納する変数。2019で割る場合、余りは0~2018までの2019通りあるので、リストの長さは2019
cnts[0] = 1 # ※1
d = 1 # sum_of_digitsの計算に使う10の累乗値

# まず余りの値ごとに余りが同じ値の数を数える
for c in s:
    sum_of_digits += int(c) * d
    sum_of_digits %= 2019
    d *= 10
    d %= 2019 # ※2
    cnts[sum_of_digits] += 1

# 次に余りの値ごとにその組み合わせを計算し、足し合わせる
ans = 0
for cnt in cnts:
    ans += cnt * (cnt - 1) // 2

print(ans)

※1

Bとしては0桁の値をとれることに注意します。例えば

A = 4038
B =
A-B = 4038

というケースです。
for c in sとやるとcは1桁の値から始まります。
そこで0桁の値の余りはfor文の前に計算しておきます。0桁の値を2019で割った余りは0なので、cnts[0]に1を代入します。

※2

d %= 2019で計算時間を抑えます。
この処理をしても結果に影響がないのは、直感的には以下で理解できます。

2020 * 10000 % 2019 = 2
2020 * 1924 % 2019 = 2 (10000 % 2019 = 1924)