問題の概要
1から9までの数字からなる文字列Sが与えられる。このとき、Sのi番目からj番目までの部分文字列を10進数の整数とみなしたとき、その値が2019の倍数になるような(i,j)の個数を求めてください。問題のURL
- 制約：
1 <= |S| <= 200000
解答
先にPythonによる解答を載せます。
s = input()[::-1] # 入力文字列を逆順でsに格納
counts = [0] * 2019
counts[0] = 1
num, d = 0, 1
for char in s:
num += int(char) * d
num %= 2019
d *= 10
d %= 2019
counts[num] += 1
ans = 0
for cnt in counts:
ans += cnt * (cnt - 1) // 2
print(ans) # 答えの出力
ダメな解法
この問題を見たとき、自分を含めて恐らく多くの方は、「(i,j)の組を総当たりすればよい」と考えたことと思います。これを素直に実装すると下記のようなコードになります。
s = input()
n = len(s)
cnt = 0
for i in range(n-2):
for j in range(i+3, n+1):
tmp = int(s[i:j])
if tmp % 2019 == 0:
cnt += 1
print(cnt)
このコードを提出するとTLE、実行時間制限超過となります。それもそのはずで、計算量を確認すると、N=2*10^5に対してO(N^2)です。したがって、Pythonだから通らないのではなく、アルゴリズムの設計（問題の解き方）が間違っています。
解法の説明
解説PDFを読むと「2019 は 2 でも 5 でも割り切れないので、mod2019 では 10 に逆元が存在します」と書かれていますが、あまりピンときませんでした。次のように考えてみましょう。入力文字列として、S = 31415926535を考えます。このとき、例えば4番目から8番目までの部分文字列15926が2019の倍数であるか判定したいというのが本問の主題です。特定の部分文字列が2019の倍数であるかどうかは、int(s[[i:j+1]) % 2019 == 0とすればよいのですが、これを逐一やっていては上でも述べたようにO(N^2)の計算量がかかってしまいます。そこで、この部分文字列を次のようにとらえてみます。
$$ 15926 = \frac{15926535 - 535}{10^3} $$
このように考えることの利点は、$10^3$ が $2019$ と互いに素（最大公約数が $1$）であるため、$15926$ が $2019$ の倍数であるかどうかは、上式右辺の分子に着目すればよいということです。したがって、アルゴリズムを設計する際は、与えられた文字列の右からk番目（$k=0,1,2,\dots,n$）から最右端までの部分文字列を保持しておけばよいことになります。これは$O(N)$で実行できます。分子が $2019$ の倍数かどうかは、分子を $A-B$ としたときに、$A \:\:\mathrm{mod}\:\: 2019$ と $B\:\: \mathrm{mod}\:\: 2019$ の値がそれぞれ等しければ、元の整数（上式左辺）は $2019$ の倍数であると言えます。以上を踏まえて、上に載せたコードを見ていただくと、なぜあのような記述で解が求まるのか分かると思います。
