AtCoder Beginner Contest 164 D を Python で

2020年4月30日

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

3

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


入力例 2

14282668646

出力例 2

2


入力例 3

2119

出力例 3 

0

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


解説

もちろん,以下のように問題文に沿ってそのまま愚直な二重ループで書くやり方は,\(O(N^2)\)で,\(N≃10^6\)なので,\(O(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)

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

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

具体的に見ていきましょう。入力例1を短くして 181718171 を用いて扱います。このとき,
\(a_1\,=\,181718171\),\(a_2\,=\,81718171\),…, \(a_{N\,-\,1}\,=\,71\),\(a_N\,=\,1\)になります。ここで,181718171 → 718 を取り出すことを考えると,これは
\(\frac{718171\,-\,171}{1000}\,=\,\frac{718000}{1000}\,=\,718\)
で得ることができます。
これはつまり,\(\frac{a_4\,-\,a_7}{10^{9\,-\,6}}\)を表していることになりますよね。つまり,\(a_i\,=\,S[i : N] \) (Sのi番目以降) だけを用いて,全ての考えうる整数をとることができます。 最初のコードだと2変数でループしなければならなかったのが,うまくやることで1変数だけで考えることができるようになります。

さて,ここからが本題です。先ほどの式\(\frac{a_i- a_{j+1}}{10^{N-j}}\)から得られた整数に対して,これが 2019 で割り切れるための条件は,
\(\frac{a_i- a_{j+1}}{10^{N-j}} \equiv 0(mod 2019) \Leftrightarrow a_i- a_{j+1} \equiv 0 (mod2019) \Leftrightarrow a_i\,=\,a_{j+1}\,(mod2019)\)
これはつまり,\(\frac{a_i- a_{j+1}}{10^{N-j}}\) を 2019 で割った余りが 0 であるということですが,2019 と 10 が互いに素であるから,これは\(a_i- a_{j+1}\)を2019で割った余りと同じであるということです。この式の分母の\(10^{N-j}\)は無視して考えられます。よって,\(a_i\)と\(a_{j+1}\)はそれぞれ2019で割った余りが同じであるということになります。

以上より,1~N までの全ての \(a_i\) を2019で割った余りを求めてやり,\(a_i\,=\,a_{j+1}\,(mod2019)\)となるような,(i, j)の組み合わせの数を出力してあげれば良いということになります。

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

では,\(a_i\) の考え方です。これは,S の N 番目から(後ろから)考えていくと,\(i \geq 0 \) として,
\(a_i\,=\,10^{i}\times s[-1-i]\,+\,10^{i-1}\,\times\,s[-1-i-1]\) + … + \(10^0\,\times\,s[-1]\) です。これは最初に d=1 で,\(a_1\,=\,S[-1] \times d\) として,この\(a_1\)に\(S[-2]\,\times\,d\,\times 10\)を足すことで得ることができます。このように考えると,ループの中で d を毎回 10 倍するような値に設定して,毎回対応する S[-i]×d をその前の値に足していくと,必要な \(a_i\) が得られます。また,d の値が大きくなりすぎると,コンピュータ内部での処理が大変になり,処理速度がかかりすぎてしまうので,毎回 2019 で割った余りにしておきます。最終的に \(a_i \times d\) の 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) の値の数え方についてです。これは まず0 が 2019 個入ったリスト(名前をcountとする) を用意します。そして各\(a_i\) の mod(2019) の値に対応するインデックス番号にあたる値 count[\(a_i\) % 2019] に1を足していくことで,最終的にこのリストの値の数が,各\(a_i\) のmod(2019) の値の総数に一致します。
また,最初の段階で,便宜上何もとらないという意味として\(a_0\,\equiv 0 (mod 2019)\) を導入したいので,count[0] = 1 としておきます。 こうすることで,例えば先ほどの例で 181718171 をとるには,\(a_5\) で表現できますが,これが 2019 の約数なら,\(a_5\,\equiv 0\,(mod\,2019)\) という条件から,便宜上用いた \(a_0\) を用いて,\(a_5 = a_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つをとる場合の数なので,\(_{count[i]}C_2\,=\,\frac{count[i] \times (count[i]-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 やここのコメントでお伝えください。
それでは!

他の問題も解説しているので,ぜひ見ていってください!
ABC 164 A
ABC 164 B
ABC 164 C
ABC 164 D