桁dpの世界【競技プログラミング】

# はじめに
- 過去に１度、[桁DPの問題を解いてみた](https://qiita.com/kei_tnk/items/9a7380b3fa954688f254#%E5%B0%91%E3%81%97%E9%9B%A3%E3%81%97%E3%81%84%E5%95%8F%E9%A1%8C)けど、記憶に定着させるため、もう一度くらい解いてみようと思う。
- 前回解いてみたときは、よくもまぁ、こんなこと考えつくなぁと思ったけど、他の問題も、同じような手順だったから、慣れれば何ともないのかもね。

# 今回のお題
- AtCoderの典型90問より問題5です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2733720/8cc3e53f-4219-6253-749c-108ee77474fe.png)
- 上図の通り、ゼロを除いた9種の数字{1,2,3,4,5,6,7,8,9}から、いくつかの数字を指定されて、N桁の数値を作る場合、Bの倍数はいくつ出来るか？という問題で、入力例は以下の通り。
```
10000 27 7 // 10000桁の数値、27の倍数、7種の数字
1 3 4 6 7 8 9 //7種の数字の内容
```
- とりあえず、dpを使うことは確定しているけど、どうやって、遷移式を作れば良いのか？その前にdpは何を表すこととすれば良いのか？
- でも、桁DPの場合は、悩みは不要！
    - dp[i][??]について、i桁の数値で、??を指標として表せる条件を満たす数値の数。
    - 上記を下記のようなfor文でiの値を0⇒n-1まで遷移させて、dp配列を埋めていく。
    - 桁dpは、当然、桁が増えれば増えるほど、数値の数は増加していく。前記入力のように7種の数値を使う場合は、単純に考えれば、1桁増える毎に、7倍に増えるイメージだろうか？
    - もし、何も条件のない5桁の数値であれば、10000〜99999で89999種類の数値があり、6桁の100000〜999999は899999種類の数値があるため、約10倍となっている。
- さて、今回は、N桁の数値がBの倍数となる場合の数なので、使用可能な数字の配列をcs[k](k in 0..<K)としたとき、遷移式は以下の用に作れる。
```swift
for i in 0..<N {
    for j in 0..<M {
        for k in 0..<K {
    		dp[i + 1][(10 * j + cs[k]) % B] += dp[i][j]
        }
    }
}
```
- 上記で、dp[i][j]は、i桁で、「mod_Bがjとなる」場合の数。modって何？って人は、[この投稿](https://qiita.com/kei_tnk/items/4a643b2f47d2aa919a1b#mod%E3%81%AE%E5%85%AC%E5%BC%8F)でも読んでね。
- 上のforループを見ると、i,jが固定されている間、当然dp[i][j]は固定されているけど、kのforループが回る間、dp[i+1]に対して+=がk回されるから、dp[i+1]計はdp[i]計のk倍の数になると分かる。
- この数値に、cs[k]を最後の桁に追加して新たな数字を作れば、その数値の「mod_B」の値は、当然(10 * i + cs[k])%Bとなるので、場合の数は上記のように遷移する。
    - ところで、この遷移式の左辺と右辺を結ぶ記号が「=」ではなく「+=」であることの意味を理解出来ているだろうか？理解出来ない場合は、前記「でも、桁DPの場合は、悩みは不要！」と書いた文の後の3行目･4行目のコメントを思い出してね。
- なお、jとｋのforループは入れ替え可能。遷移を示すiのforループは当然だけど、一番外側から変えられないよ。
- さて、このdpを使ってコードを書くと、次のとおり。
```swift
extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
}

let (N,B,K) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1],$0[2])}[0]
let cs = readLine()!.split(separator:" ").map{Int($0)!}
let mod = 1000000007


// dp準備
var dp = [[Int]](N+1,[Int](B,0)) // 0で初期化 -- これって重要①！！！


// 初期化
dp[0][0] = 1 // ゼロ桁の数値は0となり、0は当然Bの倍数となるので、1件  -- これって重要②!!!

// 遷移
for i in 0..<N {
	for j in 0..<B {
		for k in 0..<K {
            let j_next = (10 * j + cs[k]) % B
			dp[i + 1][j_next] += dp[i][j] // -- 遷移式
			dp[i + 1][j_next] %= mod
		}
	}
}

print(dp[N][0])
```
- 上記のコードは、重要①と重要②が意外と重要なポイントだと思う。
    - 遷移の初回は、`dp[1][j_next] += dp[0][j]`をj,kのforループで回すことになるけど、重要①②により、dp[0][0] = 1、dp[0][not 0] = 0となっている。
    - 重要②では、「ゼロ桁の数値は0となり、0は当然Bの倍数となるので、1件」みたいな、それっぽい言葉で､それらしいことを書いたけど、実際、dp[0][1],dp[0][2],...,dp[0][B-1]の値を全て、「0」ではなく「1」としたらどうなるか？
    - 全て「1」とすると、`dp[1][j_next] += dp[0][j]`の右辺は常に1となる。よって、j,kのループ「B×K」回転で毎回+1追加されるため、dp[1][j]は全てのj(in 0..<B)について、値Kとなる。1桁の数値しかなく、mod_Bがjになる数値の種類がK個もあるとか明らかに多すぎる！
    - よって、元の条件に戻って、dp[0][0] = 1、dp[0][not 0] = 0で考えてみると、遷移式より、dp[1][cs[0]%B]=1,dp[1][cs[1]%B]=1,...,dp[1][cs[k-1]%B]=1となり、それ以外のdp[1][??]=0となる。
    - これは、1桁のみのとき、そもそもcs[0],cs[1],...,cs[k-1]しか存在し得ないことと整合性がある。
    - よって、日本語でそれっぽく「ゼロ桁の数値は0となり、0は当然Bの倍数となるので、1件」と書いてしまったけど、実際に正しかったと言うことが実感できるだろうか？
    - そして、dp[0][0]以外のdp[0][1],...,dp[0][B-1]も0にすることが必要だったと理解されただろうか？よく、dpを準備するとき、[[Int]](repeating:[Int](repeating:-1,count:B),count:N+1)などと、値が未定であるという意味で「-1」で初期化することがあるけど、桁DPの場合など、0で初期化する必要があるケースもあるので注意。
- さて、解くことは出来たけど、N,B,Kの三重forループとなっているため、計算量はO(N･B･K)となり、アホみたいに大きくなります。Nが$10^{18}$の時なんて、確実にTLEです。

# 計算量の削減に取り組む
- Nが$10^{18}$の場合でも、解けるように、コードを改良します。
- そもそも、遷移式の左辺と右辺が += で結ばれており、右辺も左辺もdp項のみで、その他の加減がないため、配列dpの遷移は、B×Bの正方行列と要素数Bの列ベクトルの乗算でも表すことが出来ると、なんとなく想像できるだろうか？
- 遷移一つ分(i⇒i+1)が正方行列を1回乗じたモノとすれば、n回の遷移は正方行列をn乗した後に生成された正方行列を要素数Bの列ベクトルにかけても良い。
- よって、この正方行列のn乗計算を高速化するのが今回のミッションとなる。
- ん？n乗計算の高速化？そういえば・・・・、[ヤってんな！](https://qiita.com/kei_tnk/items/a622eec4e7a9d0c47aca)そうか〜、ここで「繰り返し二乗法」の威力が発揮されるのか〜、実は繰り返し二乗法を投稿したとき、面白い視点だけど何の役に立つの？と思ってたけど、こんな場合に役立つのね。
- それじゃ、まずは、dp方式から正方行列方式に移行してみますか。
- 文章だけで理解するのは無理があるので、B=4,cs=[2,3,4,5,7,9]の場合の遷移行列の作り方を書くよ！
![スクリーンショット 0006-08-31 9.55.51.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2733720/5b927502-b5b4-37c7-274a-bf910fafafeb.jpeg)
- お分かりいただけただろうか？
- 上記の理解に基づき、遷移行列を配列dp_matとして表すと、以下のコードになる
```swift

extension Array {
    init(_ size: Int, _ value: Element) {
        self = [Element](repeating: value, count: size)
    }
}

let (B,K) = (4,6)
let cs = [2,3,4,5,7,9]

// dp準備
var dp_mat = [[Int]](B,[Int](B,0)) // 0で初期化

for j in 0..<B {
	for k in 0..<K {
        let j_next = (10 * j + cs[k]) % B
		dp_mat[j][j_next] += 1

	}
}

for row in dp_mat {
    print(row) // 出力は下記
}

///出力(dp_mat)
[1, 2, 1, 2]
[1, 2, 1, 2]
[1, 2, 1, 2]
[1, 2, 1, 2]
```
- さて、行列の乗算はどのように書けば良い？ついさっき、[swiftで行列計算をするための構造体Matrixを自作](https://qiita.com/kei_tnk/items/eb4d2829e85a59ae8522)したから、それを利用することとする。
- 具体的には、こんな構造体
```swift
struct Matrix<Element:Numeric>:CustomStringConvertible{
    let size:Int
    var data:[[Element]]
    let mod:Int?
    
    var description : String {
        var str = ""
        for raw in data {
            str += raw.description + "\n"
        }
        str.removeLast()
        return str
    }
    
    init(_ cnt:Int, _ v:Element = Element.zero){
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:v,count:cnt),count:cnt)
        self.mod = nil
    }
    init(_ cnt:Int, _ vs:[Element]){
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:Element.zero,count:cnt),count:cnt)
        self.mod = nil
        var num = 0
        OUT:for i in 0..<cnt {
            for j in 0..<cnt {
                data[i][j] = vs[num]
                num += 1
                if num == vs.count {break OUT}
            }
        } 
    }
    init(_ cnt:Int, _ v:Element = Element.zero, mod:Int? = nil) where Element == Int {
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:v,count:cnt),count:cnt)
        self.mod = mod
    }
    init(_ cnt:Int, _ vs:[Element],mod:Int? = nil) where Element == Int {
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:Element.zero,count:cnt),count:cnt)
        self.mod = mod
        var num = 0
        var vs_mod = vs
        if let m = mod {vs_mod = vs.map{$0 % m}}
        OUT:for i in 0..<cnt {
            for j in 0..<cnt {
                data[i][j] = vs_mod[num]
                num += 1
                if num == vs.count {break OUT}
            }
        } 
    }
    
    // i行目、j列目を表示（ただし、0起点）
    subscript(_ i:Int,_ j:Int) -> Element {
        get {data[i][j]}
        set {data[i][j] = newValue}
    }

    static func *=(left:inout Self,right:Self) where Element == Int {
        let cnt = left.size
        var newMat = Self(cnt,mod:left.mod)
        for i in 0..<cnt {
            for j in 0..<cnt {
                for k in 0..<cnt {
                    newMat[i,j] += left[i,k] * right[k,j]
                    if let m = left.mod {newMat[i,j] %= m}
                }
            }
        }
        left = newMat
    }

    func pow(_ n:Int)->Self where Self == Matrix<Int> {
        var ans = self
        var x = self
        var i = n - 1
        while i > 0 {
            if i % 2 == 1 {ans *= x}
            x *= x
            i >>= 1
        }
        return ans
    }
    
    //列ベクトルへの乗算
    static func *(left:Self,right:[Element])->[Element]{
        let cnt = left.size
        var ans = [Element](repeating:Element.zero,count:cnt)
        for i in 0..<cnt {
                for j in 0..<cnt {
                    ans[i] += left[i,j] * right[j]                    
                }
        }
        return ans
    }

}
```
- 元の投稿では、加減算の演算子も定義していたけど、mod付きの累乗に必要な機能のみに絞った。あと、列ベクトルへの乗算も出来るよ。
- それじゃ、遷移行列をこのMatrixを使って書き換え、初期値となる列ベクトルも書いてみよう。nは3桁にしてみようか。
```swift
let (N,B,K) = (3,4,6)
let cs = [2,3,4,5,7,9]
let mod = 1000000007

// dp準備
var dp_mat = Matrix(B,0) // 0で初期化

// 初期値
var dp0 = [Int](repeating:0,count:B)
dp0[0] = 1

// 遷移行列作成
for j in 0..<B {
	for k in 0..<K {
        let j_next = (10 * j + cs[k]) % B
		dp_mat[j,j_next] += 1
	}
}

print(dp_mat) // 出力①
let dp_mat_N_mod = dp_mat.powm(N,mod)
print(dp_mat_N_mod) // 出力②
let dpN = dp_mat_N_mod * dp0
print(dpN) // 出力③
let ans = dpN[0]
print(ans) // 出力④

///出力
①
[1, 2, 1, 2]
[1, 2, 1, 2]
[1, 2, 1, 2]
[1, 2, 1, 2]
②
[36, 72, 36, 72]
[36, 72, 36, 72]
[36, 72, 36, 72]
[36, 72, 36, 72]
③
[36, 36, 36, 36]
④
36
```
-　うむ！凄くシンプルになった！
-　そして、気付いちゃった。初期値としての列ベクトルdp0を作成して､累乗後の遷移行列に乗じた後のdpNのインデックス[0](mod_B=0。つまり、Bで割り切れる)の値36を解答として出してみたけど...そんなことしなくても、遷移行列をN乗した後のdp_mat_N_mod[0][0]で良いよね。列ベクトルを乗じるまでもなかったよ！でもまぁ、基本は大事。
-　兎に角、これで、次みたいな$N=10^{18}$の問題にも挑めることになったね！
```
1000000000000000000 29 6
1 2 4 5 7 9
```
- 解答用のコードは、たったこれだけ！Matrix構造体のお陰だね。
```swift
let (N,B,K) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1],$0[2])}[0]
let cs = readLine()!.split(separator:" ").map{Int($0)!}

let mod = 1000000007

// dp準備
var dp_mat = Matrix(B,mod:mod)

// 遷移行列の作成
for j in 0..<B {
	for k in 0..<K {
        let j_next = (10 * j + cs[k]) % B
		dp_mat[j,j_next] += 1
	}
}

var dp_mat_N = dp_mat.pow(N)

print(dp_mat_N[0,0]) 
```
- うむ！paiza.ioでタイムアウトにならずに完了！

# B=1000への対応
- 最後の難関である。B=1000への対応がある。
- 前記コードの遷移行列の作成では二重forループが有るけど、計算量はO(B･K)だから、$10^4$程度で屁みたいな計算量。
- でも、Matrix構造体の二項演算子「+=」を見て貰うと、下記のように三重forループとなっており、left.sizeはBとなるから、B=1000のとき、計算量は$10^9$となり、TLEを引き起こす。
- よって、「\*=」の計算量を下げる工夫が必要...あれ？「*=」の高速化手法なんてないんじゃね？
- ということは、遷移行列を使うこと自体、三重forループ回避のため、やっては駄目なのね。
- でも、遷移行列を使った「繰り返し二乗法」の考え方は間違ってないと思うんだよね。
- あれ？そういえば、「繰り返し二乗法」の考え方を拡張した、「[ダブリング](https://qiita.com/kei_tnk/items/b305cebd27de68d1904a)」ってのがあったね。
- ダブリングの考えに従って、最初に考えていた、下記の遷移を書き換えてみよう。
```swift
// 遷移
for i in 0..<N {
	for j in 0..<B {
		for k in 0..<K {
            let j_next = (10 * j + cs[k]) % B
			dp[i + 1][j_next] += dp[i][j] // -- 遷移式
			dp[i + 1][j_next] %= mod
		}
	}
}
```
- 上記が時間を食うのは、Nが$10^{18}$のように非常に大きい数値になっていたからでした。
- しかし、ダブリングの考えを活用するため、
    - dp[1][?] から dp[2][??] をつくる
    - dp[2][?] から dp[4][??] をつくる
    - dp[4][?] から dp[8][??] をつくる
    - ....
- のようなことが出来れば、$10^{18}$回も面倒なループをする必要がなくなる。
- 具体的には、こんな感じに書ける
    - dp[2][(10 * j + k) % B] += dp[1][j] * dp[1][k]
    - dp[4][(100 * j + k) % B] += dp[2][j] * dp[2][k]
    - dp[8][(10000 * j + k) % B] += dp[4][j] * dp[4][k] -- ※
- 上記の※の計算イメージは、公式の解説にもあったけど、例えばこんな感じ
    - dp[4][j]に属する4桁の数字が<font color=green>3141</font>で、
    - dp[4][k]に属する4桁の数字が<font color=red>5926</font>のとき、
    - 8桁の数値<font color=green>3141</font><font color=red>5926</font>は、
    - dp[8][(10000 * j + k) % B]に属する。
- よって、「2^bi >= Nとなる最小のbi」を使って、遷移用のdp_biを新たに作ると
```swift
// dp_bi準備 
var dp_bi = [[Int]](bi+1,[Int](B,0))

// 一周目
for j in 0..<K{
	dp_bi[0][cs[j] % B] += 1
}

// 二周目以降
for i in 0..<bi {
	for j1 in 0..<B {
	for j2 in 0..<B {
        let j12_next = (10^(1<<i) * j1 + j2) % B
		dp_bi[i + 1][j12_next] += dp_bi[i][j1] * dp_bi[i][j2]
		dp_bi[i + 1][j12_next] %= mod
	}
	}
}
```
- 上記で作ったdp_biを用いて、繰り返し二乗法を使って解答を算出するのは下記コード
```swift
// 解答用の配列
var ans = [[Int]](bi+1,[Int](B,0))

// 初期値
ans[0][0] = 1

// dp_biより解答を算出
for i in 0..<bi {
    if N & (1<<i) != 0 { // ビットが立ったときのみdp_biを反映
    	for j1 in 0..<B {
    	for j2 in 0..<B {
            let j12_next = (10^(1<<i) * j1 + j2) % B
    		ans[i + 1][j12_next] += ans[i][j1] * dp_bi[i][j2]
    		ans[i + 1][j12_next] %= mod
    	}
    	}
    } else {
        for j in 0..<B{ // ビットが立ってないときは変更無し
        	ans[i + 1][j] = ans[i][j]
        }
    }
}
print(ans[bi][0])
```
- はい、ここまでで、皆さんはお気づきだと思うけど、swiftでは、$10^x$みたいな形で、累乗計算は出来ません。
- よって、`10^(1<<i)`の代わりに、配列`power10[i]`を置くこととして、power10は以下のように作ります。
```swift
// 10の累乗計算(mod B) -- mod_1000000007ではないことに注意
func modpow(a:Int,n:Int,mod:Int)->Int {
    var ans = 1
    var x = a
    var i = n
    while i > 0 {
        if i % 2 == 1 {ans *= x}
        x *= x
        i >>= 1
        
        //mod計算
        if ans > mod {ans = ans % mod}
        if x > mod {x = x % mod}
    }
    return ans
}
var power10:[Int] = []
for i in 0..<bi {
    power10.append(modpow(a:10,n:1<<i,mod:B))
}
```
- 以上を全て反映すると、解答を出すコードは
```swift
extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
}

let (N,B,K) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1],$0[2])}[0]
let cs = readLine()!.split(separator:" ").map{Int($0)!}
let mod = 1000000007

// 2^bi >= Nとなる最小のbiを求める
var (bi,i) = (0,N)
while i>0 {bi+=1;i>>=1}

// 10の累乗計算(mod B) -- mod_1000000007ではないことに注意
func modpow(a:Int,n:Int,mod:Int)->Int {
    var ans = 1
    var x = a
    var i = n
    while i > 0 {
        if i % 2 == 1 {ans *= x}
        x *= x
        i >>= 1
        
        //mod計算
        if ans > mod {ans = ans % mod}
        if x > mod {x = x % mod}
    }
    return ans
}

var power10:[Int] = []
for i in 0..<bi {
    power10.append(modpow(a:10,n:1<<i,mod:B))
}

// dp_bi準備 -- 旧dp[i_old]のi_old=1,2,4,8がdp_bi[i_new]のi_new=0,1,2,3に相当する。i_old=2^i_new
var dp_bi = [[Int]](bi+1,[Int](B,0))

//初期化 旧dp[1][i] を求める -- bi=0はN=1に相当、bi=1,2,3はn=2,4,8に相当
for j in 0..<K{
	dp_bi[0][cs[j] % B] += 1
}

// 遷移
for i in 0..<bi {
	for j1 in 0..<B {
	for j2 in 0..<B {
        let j12_next = (power10[i] * j1 + j2) % B // modpow10も(mod B)
		dp_bi[i + 1][j12_next] += dp_bi[i][j1] * dp_bi[i][j2]
		dp_bi[i + 1][j12_next] %= mod
	}
	}
}

// ダブリング解答
var ans = [[Int]](bi+1,[Int](B,0))
ans[0][0] = 1
for i in 0..<bi {
    if N & (1<<i) != 0 {
    	for j1 in 0..<B {
    	for j2 in 0..<B {
            let j12_next = (power10[i] * j1 + j2) % B
    		ans[i + 1][j12_next] += ans[i][j1] * dp_bi[i][j2]
    		ans[i + 1][j12_next] %= mod
    	}
    	}
    } else {
        for j in 0..<B{
        	ans[i + 1][j] += ans[i][j]
        }
    }
}

print(ans[bi][0])
```
- これで、計算量が$O(B^2･log(N))$になったから、B=1000でも時間内に解けるようになったはず。下記入力で実行してみる。
```
1000000000000000000 957 7
1 2 3 5 6 7 9
```
- ...あれ？paiza.ioだとtimeoutだね。でも、同じアルゴリズムの公式回答（c++）をpaiza.ioに入れるとtimeoutにならない。
- しょうがないので、AtCoderのコードテストを使用したら、1000ms程度で解答できた。もしかして、paizaの時間制限って1000msなの？ちなみに、コードテストにc++版コードをぶち込んでみたら、250ms程度で解答...4倍違うのかよ...swiftは駄目だなぁ

# おわりに
- 今回は、AtCoderの典型90問の第5問に取り組んでみたけど、正方行列の生成で結構苦労したよ。
- で、苦労したのに、結局ダブリングで簡単に解答できちゃうとか悲しすぎた。
- まぁ、着実に実力はついてきてると思う。後は、実際のコンテストで4完を目指すだけ！
- でも、典型90問って、まだ5問しか終わってないとか､先が見えなさすぎる。ボスケテ...