torus711 のアレ

主に競技プログラミングの問題について書きます

Digit DP 入門

はじめに

 動的計画法のパターンの一つで,「Digit DP」とか「桁 DP」と呼ばれているものがあります.問題によって異なってくる部分もありますが,問題によらず共通している部分がかなりあるので,その部分を中心に解説してみたいと思います.

Digit DP とは?

 「ある整数 N 以下の整数であって,ある条件を満たすものの個数を求めよ」といったタイプの数え上げ問題を解くときにしばしば使われる手法です.Digit DP は,整数の各桁に入る値を 1 桁ずつ決めていくという全探索を DP 化したものです.「ある条件」の部分はとりあえず置いておいて,「N 以下の整数を数える」というのを例にしてみようと思います.答えは当然 N + 1 ですが,これを Digit DP で求めることで Digit DP の本質的なアイデアが見えてくると思います.また,求める整数は 10 進表記であるとします.

Digit DP

 先程書いたように,Digit DP は整数の各桁の値を上位の桁から順に*1 1 つずつ決めていくという全探索がベースにあります.従って,桁に対するアクセスが簡単にできた方が都合がよいです.また,N が 64 bit 整数にすら収まらない場合もあるので,入力として与えられるのは N を表す文字列 S であるとします.また,L := |S| としておきます.
 さて,Digit DP をするにあたって最低限もたなければならない状態パラメータ(要するに DP テーブルの配列添字)が 2 つあって,それは「決めた桁数」と「N 未満フラグ(造語)」です.決めた桁数というのはよくある DP と同じく,どこまで処理をしたか・次に考える桁はどこかを表すパラメータです.未満フラグというのは,今見ている状態にまとめられている整数が N 未満であることが確定したか否かを表す Boolean です.なぜこのようなパラメータをもつ必要があるのでしょうか? それは,次の桁の値を総当りするときの事を考えると見えてきます.
 今,整数の i 桁目まで決定しているとして,次の桁に入れる値を決めようとしているとします.N に於ける,この位置に対応する桁 S_i が表す値を D とします(プログラム上での S[ i ] - '0' という値です).10 進表記の整数なので,次の桁に入る値の候補は 0 以上 9 以下の整数 d ですが,実は,無条件に全ての値を使えるわけではありません.過去に決定した桁が全て N の対応する桁の値に等しいという状況を考えてみます(例えば N = 12345 に対して 123?? みたいな状態).このとき,D < d なる d を次の桁に入れてしまうと,構成される整数は N より大きくなってしまいます.これは求めたかった「N 以下の整数」ではないので,余計なものまで数え上げてしまっています.ところが,過去に一度でも N の対応する桁より小さい値を入れていれば,次の桁にどんな値を入れたとしても N を超えることはありません.この 2 つの場合を適切に分離するためのパラメータが未満フラグです.見やすさ及び確認のために並べて書くと,

  • 過去に対応する桁より小さい値を入れたことがある → 全て(0 以上 9 以下)の値を自由に入れられる
  • 過去に対応する桁より小さい値を入れたことが無い → 0 以上 D 以下の範囲で決めなければならない

となります.
 このようにして,構成される整数が N を超えないように各桁の値を決めていくのが,Digit DP の基本的なアイデアです.

実装例

 ここで,わたしがよくやる実装方法を紹介しておきます.DP の状態の取り方を

  • dp[ 決めた桁数 ][ 未満フラグ ] := 総数

として dp[ i ][ j ] からの遷移を考えます.このとき,D = S[ i ] - '0' であるとして,j が 0 なら D 以下,そうでなければ 9 以下の値を次に使えるので,この選択肢を網羅する for ループは条件演算子を使って

for ( int d = 0; d <= ( j ? 9 : D ); ++d )

と書けます.また,d が決まると遷移先の未満フラグも j || ( d < D ) と決まります.つまり,

dp[ i + 1 ][ j || ( d < D ) ] += dp[ i ][ j ];

みたいな更新式を書くことができます.このようにすると割とすっきり書けるかなと思います.
 更新が終わったあと,dp[ L ][ 0 ] + dp[ L ][ 1 ] が答えになっています(ちなみに,求めるのが N 未満の整数の数だったら dp[ L ][ 0 ] だけを見ればよいです).

応用

 ここまで例にしてきた問題はとても自明なもので,Digit DP を使うまでもなく解けてしまうものでした.本題である「N 以下の整数のうちである条件を満たすものの個数を求める」という問題を解く場合には,状態遷移を適切に消すとか,この「ある条件」を満たすか否かを判別できるパラメータを増やすとかしてやる必要があります.
 例えば,「桁に 5 が出現しない整数を求める」場合であれば先程の for ループで d == 5 のときだけ continue するようにすれば数えられます.ちょっと変えて,「桁に 5 が出現する個数が丁度 M 個である整数を数える」みたいになると,今までに使った 5 の個数を状態パラメータとしてもつ必要が出てきます.
 よく出てくる追加パラメータとして,「Leading Zero フラグ(造語)」というのがあります.例えば上の例をちょっと変えて,「桁に 0M 個出現する整数を数える」としてみます.このとき,5 桁きっちり表示したときに 00505 という整数があったとして,この値の桁に含まれる 0 の数は,3 個ではなく 1 個とするのが適切です.これを正しく扱うためには,「過去に 0 以外の値を入れたことがあるか?」,すなわち「Leading Zero の部分を抜けたか?」という Boolean をパラメータに加える必要があります.

具体例 (ABC007D)

 では,実際に出題された問題を解いてみようと思います.取り上げる問題は,D: 禁止された数字 - AtCoder Beginner Contest #007 | AtCoder です.問題概要は,「正整数 A, B ( A \leq B \leq 10^{ 18 } ) について,A 以上 B 以下の整数の内,桁に 4 または 9 が 1 つ以上含まれているものの総数を求めよ」です.まず重要な点として,A 以上 B 以下での総数というのは,0 以上 B 以下の総数から,0 以上 A 未満での総数を引くことで求められます.従って,任意の非負整数 N に対して,0 以上 N 以下での総数を求める関数を作れば,これを BA - 1 に対して呼び出して差をとることで問題を解くことができます.
 では,整数 N 以下の整数の内,桁に 4 または 9 を 1 つ以上含むものの総数を Digit DP で数えてみます.まず,Digit DP をするにあたって,先細説明した「決定済み桁数」と「N 未満フラグ」はパラメータにする必要があります.これだけだと 49 を含むか否かが分からなくなってしまうので,「4 または 9 を含むか?」を表す Boolean を状態に加え,

  • dp[ 決定した桁数 ][ N 未満確定か ][ 4 または 9 を含むか ] := 総数

という DP をします.
 dp[ i ][ j ][ k ] からの遷移について考えます.このとき,構成途中の整数が既に N 未満であれば(上位の桁で N の対応する桁の値未満の値を入れている状態),次の桁には 0 以上 9 以下の値を自由に入れられます.そうでない場合,今見ている桁が N との大小関係に影響し得るので,N の対応する桁の値より大きな値は使うことができません.
 次の桁に使える値が分かったので,その値を総当りして更新します.今試そうとしている値を d として,その状態遷移で各パラメータ i, j, k がどのように変化するかを考えてみます.N の対応する桁の値を D とすると,

  • i : 新たに 1 桁決めるので,i + 1 になる
  • j : 元から 1 なら 1 .そうでない場合,次の桁に入れる値が D 未満なら 1 になる
  • k : 元から 1 なら 1 .そうでない場合,次の桁に入れる値が 4 または 9 なら 1 になる

となります.j, k の変化をそれぞれ 1 つの式で書くと

  • j || ( d < D )
  • k || d == 4 || d == 9

です.つまり,dp[ i ][ j ][ k ] から足し込む先は dp[ i + 1 ][ j || ( d < D ) ][ k || d == 4 || d == 9 ] です.
 この DP を回し終わったあと,dp[ L ][ 0 ][ 1 ] + dp[ L ][ 1 ][ 1 ] が求めたかった値です.
 コードの全容を示すと次のようになります.

template < typename T > inline string toString( const T &a ) { ostringstream oss; oss << a; return oss.str(); };

long long dp[ 32 ][ 2 ][ 2 ];
// dp[ 決めた桁数 ][ 未満フラグ ][ 4 または 9 を含むか ] := 総数

long long solve( const string &S )
{
	const int L = S.size();

	fill( ( long long * )dp, ( long long * )dp + sizeof( dp ) / sizeof( long long ), 0 );
	dp[0][0][0] = 1;

	for ( int i = 0; i < L; ++i )
	{
		const int D = S[i] - '0';

		for ( int j = 0; j < 2; ++j )
		{
			for ( int k = 0; k < 2; ++k )
			{
				for ( int d = 0; d <= ( j ? 9 : D ); ++d )
				{
					dp[ i + 1 ][ j || ( d < D ) ][ k || d == 4 || d== 9 ] += dp[i][j][k];
				}
			}
		}
	}

	return dp[L][0][1] + dp[L][1][1];
}

int main()
{
	long long A, B;
	cin >> A >> B;
	cout << solve( toString( B ) ) - solve( toString( A - 1 ) ) << endl;

	return 0;
}

まとめ

 ということで,Digit DP の基本的な考え方とちょっとの具体例でした.最後にもういくつか問題を挙げておくので,練習にお役立てください.

*1:下からでもできるっぽいですがわたしは書いたことが無いので上からやります