数弱でも分かる ABC164 D問題 (剰余演算)

エンジニア就活ではコーディング試験が主流になっているようですね。
自分も採用試験に備えて最近はAtCoderのコンテストに参加してアルゴリズムの勉強をしています。
競プロについての記事を書くのは初めてで、競プロ自体も初心者なのでお手柔らかにお願いします。

この記事では[ABC164 D問題](https://atcoder.jp/contests/abc164/tasks/abc164_d) および、それに伴って数字からなる文字列の剰余について考察します。
なお、使用言語はPython3です。

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

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

# 愚直にやると間に合わない

まず思いつくのは `i < j` という条件のもと、文字列`S`を「`i`文字目の前」と「`j`文字目の前」で文字列を切断して、その間の部分が2019で割り切れるかを計算する方法です。

---

#### 文字列`S`の例 `4252714`

|インデックス|0|1|2|3|4|5|6|
|---|---|---|---|---|---|---|---|
|文字列|4|2|5|2|7|1|4|

$i = 2, \ j = 6$ とすると、スライスした文字列 `S[i:j]` は $5271$

---

実際に提出してみるとわかりますが、計算量が $O(N^3)$ となるこの方法ではTLEになります。

```python:TLEになる解法
def main():
    s = input()
    t = 0
    for j in range(len(s)+1):
        for i in range(j-3):
            if int(s[i:j]) % 2019 == 0:
                t += 1
    print(t)

if __name__ == '__main__':
    main()
```

# 解説を見てもわからなかった

結局コンテストの時間内に解法がわからなかったので解説のpdfを読んだのですが、結局何をいっているのか分からない😢
自分のような悲しい境遇の人に向けてこの記事を書いています。

ネットでいろいろ調べて、以下の記事などを参考にしながらようやく解法が理解できたので、なるべく噛み砕いて解説したいと思います。

### 参考にした記事
* [ABC158【E問題の復習】今回も解説動画を何回も見て、何度も書いて、理解できたと思う](https://note.com/tak_nakamori/n/nb274605cd2fb)

実は、今回扱っている ABC164 D問題は [ABC158 E問題](https://atcoder.jp/contests/abc158/tasks/abc158_e)の類題です。(後で知りました)
それに伴って、この記事では以下のように問題を一般化して考えます。

# 今回考える問題 (ABC158 E問題と同じ)

長さが $N$ で $0\sim9$ の整数からなる文字列 $S$ と素数 $P$ が与えられます。
この文字列 $S$ の空でない連続する部分文字列を十進法の整数とみなしたときに、その整数が素数 $P$ で割り切れるものの個数を求めなさい。

# 合同式を使おう

ここからは合同式を用いて議論を進めていきます。「合同式って何だっけ？」って人は以下の記事などで合同式について見てください。知ってる人は読み飛ばしてもらって構いません。

* [高校数学の美しい物語 合同式の意味とよく使う６つの性質](https://mathtrain.jp/mod)

ここでは、このページでも紹介されている「合同式の和・差・積・商」を用います。
以下に記す合同式の性質についての証明はここでは割愛します。
今後特に指示がない限り、合同式の法は $p$ とします。

#### 合同式の和
```math
a \equiv b , c \equiv d \Longrightarrow a + c \equiv b + d
```

#### 合同式の差
```math
a \equiv b , c \equiv d \Longrightarrow a - c \equiv b - d
```
#### 合同式の積
```math
a \equiv b , c \equiv d \Longrightarrow a c \equiv b d
```
#### 合同式の商

$a$と$p$が互いに素であるならば

```math
ac \equiv ab \Longrightarrow  c \equiv b 
```

合同式は割り算だけ少し厄介で、両辺を割るときは割る数がpと互いに素でなければなりません。

# 具体例を見ながらアルゴリズムを理解する

以下では参考にした記事（[ABC158【E問題の復習】今回も解説動画を何回も見て、何度も書いて、理解できたと思う](https://note.com/tak_nakamori/n/nb274605cd2fb)）に倣って、 $N=8, \ P=7, \ S=43980152$ の場合について考えます。

## 前準備

文字列 $S$ に対して右から順に以下のようにインデックスを与えます。

| `i` | 7 | 6 | 5 | 4 | 3 | 2 | 1 | 0 |
|---|---|---|---|---|---|---|---|---|
|`S[i]`| 4|3|9|8|0|1|5|2|

そして新しい配列`A`について以下のルールに従って要素`A[i]`を定義します。

```python
A[i] = S[i] * pow(10, i) % 7
```
数式で書くと

```math
A[i] \equiv S[i] \times 10^i \ \ (\text{mod} \ 7 )
```

つまり、文字列 $S$ を十進数と見たときに各桁ごとに`7`の剰余を計算します。
例えば `i=4`のときは $80000$ を $7$ で割った余り $4$ を $A[4]$ に代入します。

この作業を $i=0$ から $i=N$ まで行うと以下の表のようになります。

| `i` | 7 | 6 | 5 | 4 | 3 | 2 | 1 | 0 |
|---|---|---|---|---|---|---|---|---|
|`S[i]`| 4|3|9|8|0|1|5|2|
|`A[i]`| 5|3|3|4|0|2|1|2|

---

次の作業です。今度は新しい配列`D`に `S[:i+1]` を7で割った余りを代入していきます。

| `i` | 7 | 6 | 5 | 4 | 3 | 2 | 1 | 0 |
|---|---|---|---|---|---|---|---|---|
|`S[i]`| 4|3|9|8|0|1|5|2|
|`A[i]`| 5|3|3|4|0|2|1|2|
|`D[i]`| |||||||2|

```python
D[i] = int(A[:i+1]) % 7
```

例えば `i=4`のときは $80152$ を $7$ で割った余りを代入しますが、
合同式の和の性質を用いることで $i > 0$ において以下のようになります。

```python
D[i] = (D[i-1] + A[i]) % 7
```

数式で表すと

```math
80152 \equiv 80000 + 152  \ \ (\text{mod} \ 7 )
```

$80000$ の剰余は $A[i]$ で、 $152$ の剰余は $D[i-1]$ で求められているものを利用します。

```math
80152 \equiv 80000 + 152 \equiv 4 + 5 \equiv 2 \ \ (\text{mod} \ 7 )
```

このようにして $D[4]$ は $2$ と求められます。

これを $i=1$ から $i=N$ まで行うと以下の表のようになります。

| `i` | 7 | 6 | 5 | 4 | 3 | 2 | 1 | 0 |
|---|---|---|---|---|---|---|---|---|
|`S[i]`| 4|3|9|8|0|1|5|2|
|`A[i]`| 5|3|3|4|0|2|1|2|
|`D[i]`| 6|1|5|2|5|5|3|2|

## ここから本題
この表を使うと $O(N+P)$ で答えが求められます。


例えば、 $i = 2,j = 5$として文字列 $S$ を切断したときの部分文字列 `980` について見ていきましょう。
$980 \div 7 \ \ (\text{mod} 7)$ を直接計算せず、表を使って求めます。

```math
980000 \equiv 980152 - 152 \ \ (\text{mod} \ 7 )
```

$980152$ と $152$ の剰余は表の $D[5]$ と $D[2]$ を見れば分かるので

```math
980000 \equiv 980152 - 152 \equiv 5 - 5 \equiv 0 \ \ (\text{mod} \ 7 )
```

すなわち

```math
980000 \equiv 0 \ \ (\text{mod} \ 7 )
```

となります。

さて、本当に知りたいのは $980 \div 7$ なので合同式の商の性質を使います。
ここで思い出して欲しいのが先に説明した「**割る数 $a$ と合同式の法 $p$ は互いに素でなければならない**」ということです。

今は合同式の法は $7$ なので $10$ で割っても問題ありませんが、 ABC158の問題を想定すると、合同式の法 $p$ が $10$ と互いに素であることを確認する必要があります。つまり、$10^n$ は素因数に $2$ と $5$ しか持たないので、合同式の両辺を $10^n$で割りたい場合は、合同式の法が $2$ や $5$ の倍数でないかを確かめれば良いのです。

以上の議論を踏まえた上で

```math
980000 \equiv 980 \equiv 0 \ \ (\text{mod} \ 7 )
```

となり、 $980$ が $7$ で割り切れることがわかります。

この話で何が言いたかったかというと、 $D[j] = D[i]$ のとき `S[i : j+1]` は7で割り切れるということです。
ここまでついて来れない人は他の $(i,j)$ で試してみましょう。

以上から、配列$D$について各要素の同じ数がいくつあるかを数えるだけで、「**十進法の整数とみなした部分文字列が素数 $P$ で割り切れるものの個数**」を求められます。

| `i` | 7 | 6 | 5 | 4 | 3 | 2 | 1 | 0 |
|---|---|---|---|---|---|---|---|---|
|`S[i]`| 4|3|9|8|0|1|5|2|
|`D[i]`| 6|1|5|2|5|5|3|2|

今まで考えてきた例の場合では、Dに同じ数値が2つ以上含まれているのは`5`と`2`で、それぞれ$3$つと$2$つあります。
したがって、`5`については $\_3C_2 = 3$ 通りの数で割り切れ、`2`については $\_2C_2 = 1$ 通りの数で割り切れます。
実際に、 $S=43980152$ の連続する部分文字列のうち、$7$で割り切れるのは「$980, 98 , 0 , 8015$」の4つだけです。

ここに関してはうまく説明できないのですが、実際の配列 $D$ は長さが $N$ではなく $N+1$ として、 $D[N+1]=0$ という空き部屋を一つ用意します。これは文字列$S$をスライスするインデックス$i,j$の選び方が$\_{N+1}C_2$通りあることと関係があります。極端な例で$N=1, \ P=3, \ S=3$だとしたら$D = [0,0]$としないと、計算がうまくいきませんよね。

以上のようにして部分文字列の剰余について考察してみました。最後に、[ABC158 E問題](https://atcoder.jp/contests/abc158/tasks/abc158_e)と[ABC164 D問題](https://atcoder.jp/contests/abc164/tasks/abc164_d)の解答例を載せておきます。

# ABC158 E問題 の解答例 

アルゴリズムは間違ってないがTLEになったので実際には配列Dを作成せずに求めました。C++なら配列Dを作る方針でもいける。

```python:TLEになった解答
from collections import Counter

def main():
    N,P = map(int,input().split())
    S = input()[::-1]
    ans = 0
    if P == 2 or P == 5:
        for i,s in enumerate(S):
            if int(s) % P == 0:
                ans += N-i
    else:
        D = [0] * (N+1)
        for i in range(N):
            D[i] =  int(S[:i+1][::-1]) % P
        C = Counter(D)
        ans = sum([c*(c-1)//2 for c in C.values()])
    print(ans)

if __name__ == '__main__':
    main()
```

長さ`P`のmodsという配列を作成し、剰余に応じて値を足していくことで、合計値を後で求める手間を省いています。

```python:ACの解答
def main():
    N,P = map(int,input().split())
    S = input()[::-1]
    ans = 0
    if P == 2 or P == 5:
        for i,s in enumerate(S):
            if int(s) % P == 0:
                ans += N - i
    else:
        mods = [0] * P
        mods[0] = 1
        current = 0
        x = 1
        for s in S:
            current = (current + x * int(s)) % P
            ans += mods[current % P]
            mods[current % P] += 1
            x = x * 10 % P
    print(ans)

if __name__ == '__main__':
    main()
```

# ABC164 D問題 の解答例 

こちらは割る数が$2019$と決まっているので、ifの分岐の必要がありません。

```python
def main():
    S = input()[::-1]
    ans = 0
    mods = [0] * 2019
    mods[0] = 1
    current = 0
    x = 1
    for s in S:
        current = (current + x * int(s)) % 2019
        ans += mods[current % 2019]
        mods[current % 2019] += 1
        x = x * 10 % 2019
    print(ans)

if __name__ == '__main__':
    main()
```


# おわりに

競プロはアルゴリズムだけでなく数学の知識も必要ですね。わからない問題に出会うたびにひとつひとつ攻略していきたいと思います。
ちなみに、就活のコーディング試験では「あ、これ進研ゼミでやったやつ！」レベルで一度経験したことがある問題が出題されたことがありました。エンジニア就職目指す人は競プロやって損はないでしょう。それでは！