桁dpの初歩【競技プログラミング】

はじめに

• 過去、桁dpを取り上げたことがあったけど、頭に全く残ってなかったし、今読み返しても、頭にスッと入ってこない。
• やっぱり、初歩を飛ばしてやるからこんなことになるんだよね。
• ということで、一番基本的な問題を解いてみることとした。
• というか、基本的な問題に出会って、過去にやったことを全く覚えてなかったことに驚愕して、慌てて初歩に手を付けた、と言うのが正確な表現だね。

桁dpの一番基本的と思われる問題

• コレです。
  • 数字、KとDが与えられて、1〜K(＜$10^{10000}$)の数字のうち、各桁数字の合計がDの倍数となる数字はいくつあるか？(mod 1000000007)
• 問題だけ見ると簡単よね。計算量を考えなければ、簡単に実現できそう。

解き方

• うろ覚えだけど、以前解いたときは、各桁の数字として0〜9の全ての数字が使えるわけではなく、例えば、[0,3,5,7,8]しか使えない制約ががあったなぁ。
• 今回は、その制約がないことが大きな違い。制約がないから、簡単なはず！
• そして、前回は、「その整数自身がBの倍数か」で判断したけど、今回は「各桁の数値の合計がBの倍数か」だから、かなり違うね。
• で、前回は、桁dpとして、dp[i][k]とは、(下位)i桁以下の数字で(mod D)がkになる「場合の数」としていたね。
• だから最終的な答えとしては、dp[N][0]を求めれば良いよね。よし、頑張って思い出しながら、コードを書いてみよう
• あれ？以前解いたときの元の数値は、N桁の数字（例：5桁の場合、10000〜99999）ってだけの指定だったけど、今回の問題は、K=54321（この場合、1〜54321の範囲）のように、各桁毎で9まで使えるとは限らない！
  • 例えば、K=54321のときにdpを行う場合、頭の3桁が543の場合、4桁目は2以下にする必要があるけど、頭の3桁が542以下であれば、4桁目は0〜9の好きな数字が使える。
• よって、今回の桁dpでは、dp[i][d][0/1]として3次元dpにします
  • 元のインデックスがkだった2次元目部分は今回の条件で使われる大文字のＫと紛らわしいので倍数Dから連想されるdに変更。
  • 最後の箱(3次元目部分)は、K=54321の例であれば、上位3桁が543の時は0、542以下の時は1とするフラグです。初歩なのに、dpが3次元になったよ！
• あと、Kを分解して頭からの各桁に格納される数値をK[i]としておくのが良いみたい。ん？dp[i][d][0/1]のiって、前回解いたときは下の桁からの数だったけど、今回は上からの桁数にした方が良いのか...。まぁ、前回（その数自身）と今回（その数の各桁の合計）では、Bの倍数かどうかの判定対象が違うから、今回はiを桁の上位から数えるのは当然かな。
• なんだか、全然別物になっちゃってるね。とりあえず、コードを書いてみるよ。
• 改めて説明するけど、dp[i][d][0/1]については、N桁の数値の上位「i-1」までの数値を使ったときの余剰がdとなる組み合わせの数だよ。[0/1]については、さっきの説明の通り。
• だから、dpは0で初期化した後、初期値として、dp[0][0][0] = 1として、それを育てていくイメージだね。

extension Array { // 3次元配列初期化を短縮表記するためののエクステンション
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
    init<T>(_ cnt2:Int, _ cnt1:Int, _ v:T) where Element == [T] {
        self = [[T]](cnt2,[T](cnt1,v))
    }
    init<T>(_ cnt2:Int, _ cnt1:Int , _ cnt0:Int, _ v:T) where Element == [[T]] {
        self = [[[T]]](cnt2,[[T]](cnt1,[T](cnt0,v)))
    }
}

let K = Array(readLine()!).map{Int(String($0))!}
let D = Int(readLine()!)!
let N = K.count

let mod = 1000000007
// print(K)

var dp = [[[Int]]](N+1,D,2,0) // 短縮表記
dp[0][0][0] = 1 

for i in 0..<N {
    let k = K[i]
    for d in 0..<D{
        for t in 0..<10 {
            if t == k {
                dp[i+1][(d+t)%D][0] += dp[i][d][0] % mod
            }

            if t < k {
                dp[i+1][(d+t)%D][1] += dp[i][d][0] % mod
            } 
            
            dp[i+1][(d+t)%D][1] += dp[i][d][1] % mod
            
        }
    }
}

print((dp[N][0][0] + dp[N][0][1] - 1) % mod)

• 出来た!提出したら、300ms程度でACでした。
• まぁ、この問題のキモは、dpの3次元目の[0/1]の設定を閃くかどうかだけだから、その部分をコッソリ見ておけば、コードは書けるよね。
• ただし、提出前検証まで、気づかない点がありました。
• 最後の行の「dp[N][0][0] + dp[N][0][1] - 1」を見てください。1を引いてるよね？
• これは、K=54321の例であれば、各桁合計がBの倍数となるパターンの中に「00000」が入っちゃうから、除く必要があるため。

おわりに

• 桁dpの初心者版を解いたけど、三次元dp[i][d][0/1]の最後の[0/1]は、普通は閃かないって。これが閃かなくても、他に無理矢理解く方法はあるのかもだけど、天才じゃない人には厳しいね。
• 初心者向けだけど、暗記問題だね、これ。
• まぁ、とりあえず、こっちの桁dpの方は覚えられたと思うよ。前回の桁dpは、また忘れそうだけどね。