はじめに
- 過去に１度、桁DPの問題を解いてみたけど、記憶に定着させるため、もう一度くらい解いてみようと思う。
- 前回解いてみたときは、よくもまぁ、こんなこと考えつくなぁと思ったけど、他の問題も、同じような手順だったから、慣れれば何ともないのかもね。
今回のお題
- AtCoderの典型90問より問題5です。
- 上図の通り、ゼロを除いた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)としたとき、遷移式は以下の用に作れる。
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って何？って人は、この投稿でも読んでね。
- 上の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を使ってコードを書くと、次のとおり。
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]などと、値が未定であるという意味で「-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乗計算の高速化？そういえば・・・・、ヤってんな！そうか〜、ここで「繰り返し二乗法」の威力が発揮されるのか〜、実は繰り返し二乗法を投稿したとき、面白い視点だけど何の役に立つの？と思ってたけど、こんな場合に役立つのね。
- それじゃ、まずは、dp方式から正方行列方式に移行してみますか。
- 文章だけで理解するのは無理があるので、B=4,cs=[2,3,4,5,7,9]の場合の遷移行列の作り方を書くよ！
- お分かりいただけただろうか？
- 上記の理解に基づき、遷移行列を配列dp_matとして表すと、以下のコードになる
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を自作したから、それを利用することとする。
- 具体的には、こんな構造体
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桁にしてみようか。
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の値36を解答として出してみたけど...そんなことしなくても、遷移行列をN乗した後のdp_mat_N_mod[0][0]で良いよね。列ベクトルを乗じるまでもなかったよ！でもまぁ、基本は大事。
- 兎に角、これで、次みたいな$N=10^{18}$の問題にも挑めることになったね！
1000000000000000000 29 6
1 2 4 5 7 9
- 解答用のコードは、たったこれだけ！Matrix構造体のお陰だね。
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ループ回避のため、やっては駄目なのね。
- でも、遷移行列を使った「繰り返し二乗法」の考え方は間違ってないと思うんだよね。
- あれ？そういえば、「繰り返し二乗法」の考え方を拡張した、「ダブリング」ってのがあったね。
- ダブリングの考えに従って、最初に考えていた、下記の遷移を書き換えてみよう。
// 遷移
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桁の数字が3141で、
- dp[4][k]に属する4桁の数字が5926のとき、
- 8桁の数値31415926は、
- dp[8][(10000 * j + k) % B]に属する。
- よって、「2^bi >= Nとなる最小のbi」を使って、遷移用のdp_biを新たに作ると
// 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を用いて、繰り返し二乗法を使って解答を算出するのは下記コード
// 解答用の配列
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は以下のように作ります。
// 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))
}
- 以上を全て反映すると、解答を出すコードは
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問しか終わってないとか､先が見えなさすぎる。ボスケテ...