Re:桁DPのこと - Luzhiled's memo

この記事はICT Advent Calendar 2017 3日目の記事として書かれています

去年書いた桁DPの記事が未だに結構参照されているのですが、あまり良くない内容だなと自分でも思っているので書き直します

なお、深夜(現在05:05)に書かれているため後日加筆修正があります。
[2018/1/27 未明]　うしさんから面白い話聞いたので追記、ついでに問題を微妙に変えた

桁DPって

こういう問題知らないと解けないけど知ってたらやるだけだから知ってて損はないよ(ってDEGwerさんが言ってました)。

一言で説明するのは難しいので例題を追いながら説明していきます。

例題

$N$ 以下の自然数のうち、各桁の数の和が $D$ であるような数の個数を $mod 1,000,000,007$ で求めてください。

制約
$1 \leq N < 10^{1000}$
$1 \leq D \leq 1000$

桁DPのわかりやすいところその1

制約がクソデカい(ことが多い)。*1

計算量的に $O(lgN)$ もしくは $O( 1 )$ くらいじゃないと間に合わないので、ここまで制約が大きいと必然的に解法は限られてくると思います。
DEGwerさんが言っていたのは、そういう時に桁DPの存在を知っていると解ける可能性が非常に高くなるよって意味なんじゃないかなって思っています。違うかも。

多倍長演算でも実装されていない限りここまで大きな数は扱えないので、入力はstring型にします。*2
以下ではC++でコードを記述しています。他の言語を使用する際は適宜読み替えてください。

#include <bits/stdc++.h>

// グローバル変数にしているのはこのあと他の関数でも使えるようにするためです
string N;
int D;

int main(int argc, char *argv[]) {
  cin >> N >> D;

  return 0;
}

問題を解いていると追々わかるかと思いますが、桁DPはDPの次元が非常に多いです。
forで実装すると一部の例外を除いて頭こわれるためメモ化再帰をオススメします。*3

C/C++にはデフォルト引数*4なんて便利なものがあるので、先に出力を書いてしまいましょう。

#include <bits/stdc++.h>
using namespace std;

string N;
int D;

// メモ化再帰用の関数
int rec() {

}

int main(int argc, char *argv[]) {
  cin >> N >> D;

  cout << rec() << endl;
  return 0;
}

※ これ以降 main 関数に触れることがほとんどないので、しばらくはメモ化再帰用の関数(rec)だけを書いていきます。

少し実装からは脱線しますが、桁DPをする際の状態の持ち方 (≒DPテーブルの持ち方) について考えます。

$N = 9999999, D = 11$ だとして、以下に示す2つの7桁の数を例に考えてみましょう。
***に入る部分はまだ値が決まっていないものとします。

1134***
1242***

この時点での2つの数の各桁の数の和はどちらも9です。

$D = 11$ なので、残りの***にあてはまる3桁の数字の各桁の数の和が2になればいいわけことがわかります。
3桁の数*5のうち、各桁の数の和が2になるものを列挙してみましょう。

002
011
020
101
110
200

以上の7パターンしかないことがわかりました。

さて、一番最初に示した

1134***
1242***

ですが、この数の***に当てはまる数は先に挙げた7パターンしかないことがわかっています。
列挙する時に使った情報は

残り3桁であること (⇔4桁使っていること)
今まで見てきた桁の数の和が9であること

だけなので、この2つの状態から条件を満たす数になるパターンはどちらも7パターンとなり、同一視してしまって問題ないことがわかります。

1134***
1242***
の2つに限らず、0234***や8100***なども同一視してしまって問題ないことがわかるかと思います。

桁DPの典型的な状態の持ち方の一つが、今まで何桁見てきたかです。

今回の場合はこれまで見てきた各桁の数の和も使いましたが、このあたりは問題によって異なってきます。
数を同一視するためには最低限どんな情報が必要なのか考えることが個人的には大事だと思います。

もう一つの典型的な状態の持ち方が "今見ている桁の自由度"です。

この問題のもう一つの条件として " $N$ 以下の自然数のうち" というものがあります。
この条件について考えてみましょう。

上から桁を決めていった時、数は

分類1 : これ以降どんな数にしても条件を満たす
分類2 : 次見る桁の数までは条件を満たすことがある
分類3 : これ以降どんな数にしても条件を満たさない

の3つに分類できることになります(この辺説明がわかりにくくてすみません><)

例を挙げます。

$N = 214517$ 、上から2桁が21であり、今から3桁目を決めるとき

$210$ *** 分類1
$211$ *** 分類1
$212$ *** 分類1
$213$ *** 分類1
$214$ *** 分類2
$215$ *** 分類3
$216$ *** 分類3
......

分類1の場合、4桁目の自由度は 0 ~ 9
分類2の場合、4桁目の自由度は 0 ~ 5(Nの上から4桁目と同じ)
分類3は条件を満たすことがないためそもそも遷移しないようにする。

最後の条件は考えなくてもいいようにするため、この条件は2値(bool型)です。

これが先に言った "今見ている桁の自由度" であり、桁DPにおいて典型的な状態の持ち方の一つです。

この問題を解くために必要な考え方は以上です、あとはコードを書いていきましょう！

先に挙げた

今見ているのは先頭から何桁目か
今見ている桁の自由度
今まで見てきた桁の数の和

の3つを表す変数がこのDPの状態を表すものです。

rec()の引数はこの順番で書かれています。

int rec(int k = 0, bool tight = true, int sum = 0) {

}

デフォルト引数を忘れないようにしましょう。ぼくは今忘れてました。

さて、遷移先を決める上で桁の自由度というものが重要になってきます。まずはそこから書いてみましょう。

int rec(int k = 0, bool tight = true, int sum = 0) {

  // 今見ている桁の数
  int x = N[k] - '0';

  // tight が true なら x までしか遷移しないようにする
  int r = (tight ? x : 9);

  int res;
  res = 0;

  // r ”まで” 遷移することに注意！ 未満ではありません
  for (int i = 0; i <= r; ++i) {
    // ここに遷移を書く
  }

  return res;
}

おおよその雰囲気は伝わるでしょうか。

次は遷移の式を書いていきます。

次の状態を考えた時、
k → k + 1
sum → sum + i
はすぐに分かるかと思います。

次の桁に行く、今決めた値を各桁の和に追加する、といったイメージです

tightは少しだけややこしいですが、

tight が false の状態から true に遷移することはない
i != r のとき false

の2つが条件になっていることが少し考えるとわかるかと思います。

それではこの遷移の式をそのままコードにしてみましょう。

int rec(int k = 0, bool tight = true, int sum = 0) {

  int x = N[k] - '0';
  int r = (tight ? x : 9);

  int res;
  res = 0;

  for (int i = 0; i <= r; ++i) {
    res += rec(k + 1, tight && i == r, sum + i);
    res %= mod;
  }

  return res;
}

最後に、再帰を書く上で絶対に忘れてはいけない終了条件を書きます。

k == N.size() のとき(≒最後の桁まで見たとき)、これ以上遷移することはないのでこれが終了条件です。
このとき、sum == D であれば条件を満たしている数であるため1を、そうでない場合は0を返します。
また、k == N.size() に限らず sum > Dである場合、これ以降どんな遷移をしても解は0になるため、その条件も一緒に書きます*6

int rec(int k = 0, bool tight = true, int sum = 0) {
  if (sum > D) {
    return 0;
  }
  if (k == N.size()) {
    return sum == D;
  }

  int x = N[k] - '0';
  int r = (tight ? x : 9);

  int res;
  res = 0;

  for (int i = 0; i <= r; ++i) {
    res += rec(k + 1, tight && i == r, sum + i);
    res %= mod;
  }

  return res;
}

これでほぼ完成です！
動作の確認をしてみましょう。

sample input

100 3

sample output

100 以下の自然数のうち、各桁の数の和が3になる数は $3, 12, 21, 30$ だけなので(多分)、再帰は問題なく動作してそうです(多分)。

あとはメモ化しておわり！w

これで $O(DlgN)$ でこの問題を解くことができました。

メモ化には参照渡しという文明の利器を使いましょう。楽です。
また、dpテーブルを-1で初期化しておくことで、bit反転一つで条件分岐ができます。(-1はすべてのbitが1になります。逆に言えばすべてのbitを反転させて0(10進数)になる数は-1だけです)

以下全体のコード

#include <bits/stdc++.h>
using namespace std;

string N;
int D;

int dp[1000][2][1001];

int rec(int k = 0, bool tight = true, int sum = 0) {
  if (sum > D) {
    return 0;
  }
  if (k == N.size()) {
    return sum == D;
  }

  int x = N[k] - '0';
  int r = (tight ? x : 9);

  int &res = dp[k][tight][sum];
  if (~res) return res;
  res = 0;

  for (int i = 0; i <= r; ++i) {
    res += rec(k + 1, tight && i == r, sum + i);
    res %= mod;
  }

  return res;
}

int main(int argc, char *argv[]) {
  cin >> N >> D;

  memset(dp, -1, sizeof(dp));
  cout << rec() << endl;
  return 0;
}

DPの状態は先に挙げたもの以外にももちろん沢山出てくるかと思いますが、基本さえ押さえていれば多分なんとかなります

〜〜〜ここから追記[2018/01/27]〜〜〜

$O(f(x) * M)$ *7 の桁DPに対し、 $Q$ 個のクエリが与えられる問題を $O( (f(x) + Q) * M)$ *8 で処理するテクをうしさんから学んだので書きます

例題

$Q$ 個の自然数 $N_i (1 \leq i \leq Q)$ が与えられます。各 $N_i (1 \leq i \leq Q)$ について以下のクエリに答えてください。

$N_i$ 以下の自然数のうち、各桁の数の和が $D$ であるような数の個数を $mod 1,000,000,007$ で求めてください。

制約
$1 \leq Q \leq 1000$
$1 \leq N_i < 10^{1000}$
$i \leq D \leq 1000$

愚直にやると $O(QD lgN)$ となり、微妙に間に合わない雰囲気を感じます

先ほど書いたDPは、
$dp[何桁目を見ているか][tightか][合計] := 値$
となっています。

各クエリにおいて差が出てくる部分は $[tightか]$ の部分だけなので、各クエリにおいて tight な状態のときだけ再計算することで $O((D + Q) lgN)$ に抑えることができます。

#include <bits/stdc++.h>
using namespace std;

int Q;
string N;
int D;

int dp[1000][1001];

int rec(int k = 0, bool tight = true, int sum = 0) {
  if (sum > D) {
    return 0;
  }
  if (k == N.size()) {
    return sum == D;
  }

  int x = N[k] - '0';
  int r = (tight ? x : 9);

  if (!tight && ~dp[k][sum]) return dp[k][sum];
  int res = 0;

  for (int i = 0; i <= r; ++i) {
    res += rec(k + 1, tight && i == r, sum + i);
    res %= mod;
  }

  if (!tight) dp[k][sum] = res;
  return res;
}

int main(int argc, char *argv[]) {
  cin >> Q >> D;

  memset(dp, -1, sizeof(dp));

  while (Q--) {
    cin >> N;
    N = string(1000 - N.size(), '0') + N;
    cout << rec() << endl;
  }
  return 0;
}