AtCoder Beginner Contest 164 D を Python で

AtCoder Beginner Contest 164 D を Python で解説していきます。
この問題はコンテスト中かなり頭を悩ませたのですが，どうやら ABC 158 E 問題にとても似ている問題だったようですね。とても数学的要素が強い問題かと思いますが，数学が得意でない人や灰コーダー向けに重要な点を抜き出して解説していくので，頑張って理解していきましょう！
それでは，まずは問題文からいきます。

D – Multiple of 2019

実行時間制限: 2 sec / メモリ制限: 1024 MB

配点 : 400 点

問題文

1 から 9 までの数字のみからなる文字列 S が与えられます。

次のような条件を満たす整数の組 (i,j) (1≤i≤j≤|S|) の総数を求めてください。

条件: S の i 文字目から j 文字目までを 10 進法の整数としてみると、この整数は 2019 の倍数である。

制約

1≤|S|≤200000
S は 1 から 9 までの数字のみからなる文字列

入力

入力は以下の形式で標準入力から与えられる。

S

出力

条件を満たす整数の組 (i,j) (1≤i≤j≤|S|)の総数を出力せよ。

入力例 1

1817181712114

出力例 1

条件を満たすのは (1,5),(5,9),(9,13) の 3 個です。

入力例 2

14282668646

出力例 2

入力例 3

2119

出力例 3

条件を満たす整数の組は存在しません。

解説

もちろん，以下のように問題文に沿ってそのまま愚直な二重ループで書くやり方は，O(N^2)で，𝑁≃10^6
なので，𝑂(10^12)
程度のオーダーになって，TLE になります。
(オーダーの感覚は，大体10^8以下となることを目安にプログラムすれば間に合います。)
また，詳しく言うと，S は最大 200000 桁ですが，コンピュータ内では64 bit の整数ですら 18 桁くらいしか表現できないので（2^(63-1)という数字になります）普通の整数型では表せず，以下のような単純な実装ではかなり重い計算になってしまうようです。(Python では内部処理で無理やり数を出すこと自体はできます)もはやオーダーがどうこうという話ではなくなります。😱

s = input()
n = len(s)

count = 0

for i in range(n-3):
    for j in range(i+3, n):
        number = int(s[i:j+1])
        if number % 2019 == 0:
            count += 1

print(count)

ということで，何らかの工夫が必要になるわけです。

まず，このような区間の数え上げは，片方を固定して，もう片方を高速で列挙するのが定石です。𝑎_i=𝑆[𝑖:𝑁]
(Sのi番目以降) とします。すると，S の i 番目から j 番目までの数をこれで表そうとすると，(a_i-a_(j+1))/(10^(N-j)) と表すことができます。

具体的に見ていきましょう。入力例１を短くして 181718171 を用いて扱います。このとき，𝑎1=181718171
，𝑎2=81718171
，…, 𝑎𝑁−1=71
，𝑎𝑁=1
になります。ここで，181718171 → 718 を取り出すことを考えると，これは

(718171−171)/1000=718000/1000=718

で得ることができます。
これはつまり、(a_4-a_7)/10^(9-6)を表していることになりますよね。つまり，𝑎𝑖=𝑆[𝑖:𝑁]
(Sのi番目以降) だけを用いて，全ての考えうる整数をとることができます。最初のコードだと２変数でループしなければならなかったのが，うまくやることで１変数だけで考えることができるようになります。

さて，ここからが本題です。先ほどの式𝑎𝑖−𝑎𝑗+1/10^(N-j) から得られた整数に対して，これが 2019 で割り切れるための条件は，(a_i-a_j+1)/(10^(N-j))≡0(𝑚𝑜𝑑2019)⇔𝑎𝑖−𝑎𝑗+1≡0(𝑚𝑜𝑑2019)⇔𝑎𝑖=𝑎𝑗+1(𝑚𝑜𝑑2019)
これはつまり，𝑎𝑖−𝑎𝑗+1/10^(N-j) を 2019 で割った余りが 0 であるということですが，2019 と 10 が互いに素であるから，これは𝑎𝑖−𝑎𝑗+1を2019で割った余りと同じであるということです。この式の分母の10𝑁−𝑗
は無視して考えられます。よって，𝑎𝑖
と𝑎𝑗+1
はそれぞれ2019で割った余りが同じであるということになります。

以上より，１～N までの全ての 𝑎𝑖 を2019で割った余りを求めてやり，𝑎𝑖=𝑎𝑗+1(𝑚𝑜𝑑2019)となるような，(i, j)の組み合わせの数を出力してあげれば良いということになります。

そのために，これまで𝑖 は 1から順に N までという順番で考えていましたが，逆に N から 1 までの順番で考えてあげます。この順に，𝑎𝑖の値をリストに入れていくことによって，先ほど考えていた𝑗をこの𝑖で考えていることになり，𝑎𝑖+1 で取得したい場合の数は，𝑖 番目までにリストに入れていたmod(2019) の値と同じものの数を順次考えていくだけで算出することができ，より効率よく数え上げることができます。

では，𝑎𝑖の考え方です。これは，S の N 番目から（後ろから）考えていくと，𝑖≥0
として，𝑎𝑖=10^i×𝑠[−1−𝑖]+10^𝑖−1×𝑠[−1−𝑖−1]+ … +10^0×𝑠[−1]
です。これは最初に d=1 で，𝑎1=𝑆[−1]×𝑑
として，この𝑎1
に𝑆[−2]×𝑑×10
を足すことで得ることができます。このように考えると，ループの中で d を毎回 10 倍するような値に設定して，毎回対応する S[-i]×d をその前の値に足していくと，必要な 𝑎𝑖 が得られます。また，d の値が大きくなりすぎると，コンピュータ内部での処理が大変になり，処理速度がかかりすぎてしまうので，毎回 2019 で割った余りにしておきます。最終的に𝑎𝑖×𝑑
の mod(2019) の値を出せばよいので，先にd を 2019 で割っても (mod 2019) の値は変わりません。
ここまでは以下のようなコードで表せます。

# a_iはこのnumber*dで表していく
number = 0
d = 1

for i in s:
    # a_iを表現
    number += int(i)*d
    # numberで次の桁にしたいから，dを10倍
    d *= 10
    # numberが大きすぎない値になるようにd自体を先にmod2019の値でとる
    d %= 2019

次に，mod(2019) の値の数え方についてです。これはまず０が 2019 個入ったリスト（名前をcountとする) を用意します。そして各𝑎𝑖の mod(2019) の値に対応するインデックス番号にあたる値 count[𝑎𝑖% 2019] に１を足していくことで，最終的にこのリストの値の数が，各𝑎𝑖のmod(2019) の値の総数に一致します。
また，最初の段階で，便宜上何もとらないという意味として𝑎0≡0(𝑚𝑜𝑑2019)
を導入したいので，count[0] = 1 としておきます。こうすることで，例えば先ほどの例で 181718171 をとるには，𝑎5 で表現できますが，これが 2019 の約数なら，𝑎5≡0(𝑚𝑜𝑑2019)
という条件から，便宜上用いた𝑎0 を用いて，𝑎5=𝑎0として，場合の数を計算できます。

ここまでを書くと，以下のようになります。

s = input()
# sを逆にする
s = s[::-1]
n = len(s)

# 各インデックス番号にmod2019の値が格納されるリストを設定
count = [0]*2019
# a_0の分を入れておく
count[0] = 1
# a_iはこのnumber*dで表していく
number = 0
d = 1

for i in s:
    # a_iを表現
    number += int(i)*d
    # a_iのmod2019の値をcountに格納
    count[number%2019] += 1
    # numberで次の桁にしたいから，dを10倍
    d *= 10
    # numberが大きすぎない値になるようにd自体を先にmod2019の値でとる
    d %= 2019

後はこの count の中に格納された値について，取りうる場合の数の和を出します。これは，count のそれぞれのインデックス番号に対する値について，異なる２つをとる場合の数なので，𝑐𝑜𝑢𝑛𝑡[𝑖]𝐶2=𝑐𝑜𝑢𝑛𝑡[𝑖]×(𝑐𝑜𝑢𝑛𝑡[𝑖]−1)/2を，count 内の全ての値について取った和を出して，その値を出力します。

以上，詳しく説明しすぎて逆によく分かりにくくなってしまった感がありますが，いかがだったでしょうか。😅
それでは，解答例に移ります。

解答例

s = input()
s = s[::-1]
n = len(s)

count = [0]*2019
count[0] = 1
number = 0
d = 1

for i in s:
    number += int(i)*d
    count[number%2019] += 1
    d *= 10
    d %= 2019

ans = 0

for i in count:
    ans += i*(i-1)//2

print(ans)

個人的に，今までで一番説明が難しく，冗長になってしまったような気がしますが，どうしょうか。ABC 158 E とほぼ同じだったようで，自分もこの回のコンテストには参加していたので，解けなかったのが悔しかったです。もし分からないところや，間違っているところがあれば，twitter やここのコメントでお伝えください。
それでは！

タムログ

プログラミングについての情報発信をしていきます！