mod を用いた計算問題

初めに

AtCoder で以下のような問題が出題された際に全く手が出なかった為、理論とその実装についてまとめてみる。

Multiple of 2019 ( AtCoder Beginner Contest 164; Problem D )

問題文
1 から 9 までの数字のみからなる文字列 S が与えられます。
次のような条件を満たす整数の組 (i,j) (1≤i≤j≤|S|) の総数を求めてください。

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

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

理論

愚直に for ループを何重にも重ねれば計算は可能であるが、AtCoder の場合は計算時間の上限を超えてしまう為に工夫が必要である。

この手の倍数を計算する問題（もしくは進数を考える問題も？）は mod の計算を上手く利用したアルゴリズムにより効率的なプログラムを書くことが出来る。

mod の基本公式

a \equiv b \pmod{ 2019 } \\
c \equiv d \pmod{ 2019 }

のとき

a + c \equiv b + d \pmod{ 2019 } \\
a - c \equiv b - d \pmod{ 2019 } \\
a * c \equiv b * d \pmod{ 2019 }

が成立する。つまり mod の計算にも通常の四則演算における和・差・積の公式が適用できるということである。

因みに商については、以下のルールも成り立つ。

$ab \equiv ac$ でかつ$a$と$p$が互いに素なら、$b \equiv c$

本問題への適用

与えられた文字列$S$が

S = s_n s_{n-1} \dots s_2 s_1

であるとする。つまり問題とは逆で$i<j$とする設定で考えていく。問題における条件を満たす組$(i,j)$とは

N_{ij} = 10^{i-j} s_i + 10^{i-j-1} s_{i+1} + \dots + 10 s_{j+1} + s_j

が2019の倍数であるもののことである。この条件の組$(i,j)$を数えることが求められている。
まず以下の$r_i$と$t_i$を定義する。

t_i = 10^{i-1} \\
r_i = t_i s_i + t_{i-1} s_{i-1} + \dots + t_1 s_1

この$r_i$は以下のような書き方をすることも出来る。

r_i = t_i s_i + r_{i-1}

このとき以下の式が成り立つ。

---

r_{i} \equiv r_{j-1} \Rightarrow N_{ij} \equiv 0 \pmod{ 2019 } 

---

以下、この証明を示す。$r_i$と$N_{ij}$の間には次のような関係が成り立っている。

r_{i} - r_{j-1} = N_{ij} * 10^{j-1} 

もし、$r_{i} \equiv r_{j-1}$である場合

r_{i} - r_{j-1} = N_{ij} * 10^{j-1} \equiv 0 \pmod{ 2019 } 

となる為、

N_{ij} \equiv 0 \pmod{ 2019 } 

または

10^{j-1} \equiv 0 \pmod{ 2019 } 

が成立するはずである。しかし $10^{j-1} = 2^{j-1} * 5^{j-1}$ であるので、$2019$ と $10^{j-1}$ は明らかに互いに素であり、前者の $N_{ij} \equiv 0$ が成り立つことになる。

アルゴリズム

r_i \equiv R_i \pmod{ 2019 } \\
s_i \equiv S_i \pmod{ 2019 } \\
t_i \equiv T_i \pmod{ 2019 }

ただし、$R_i,S_i,T_i$ は余りを表しており $0 \le R_i,S_i,T_i \le 2018$ である。
これらの文字を用いると各 $r_i$ の余りは次のように連続的に計算できる。

\begin{align}
 r_i &= t_i s_i + r_{i-1} \\
& \equiv T_i S_i + R_{i-1} = R_i \pmod{ 2019 } \\
\end{align}

この方法を用いて各 $r_i$ の2019に対する余りを計算し、余りが等しい $r_i$ の組み合わせを数え上げることで答えを計算することが出来る。

実装

上記のアルゴリズムを用いることで、for ループを重ねずに計算することが出来る。

s = input()

d = [ int( 0 ) ] * 2019
d[ 0 ] = int( 1 )
r = int( 0 )
t = int( 1 )

for si in reversed( s ):
#    print( si )

    r = int( si ) * t + r
    r = r % 2019

    t = t * 10
    t = t % 2019

    d[ r ] = d[ r ] + 1

ans = int( 0 )
for di in d:
    ans = ans + di * ( di - 1 ) // 2

print( ans )