AtCoder Beginner Contest 184, D : increment of coins

問題文

　https://atcoder.jp/contests/abc184/tasks/abc184_d

問題概要

　金貨が $A$ 枚，銀貨が $B$ 枚，銅貨が $C$ 枚ある．いずれかの種類のコインが $100$ 枚になるまで以下の試行を繰り返す．

1. コインを一様ランダムに選ぶ
2. 選ばれたコインと同種のコインを $1$ 枚追加する

操作回数の期待値はいくらか？

制約

• $0 \leq A, B, C \leq 99$
• $0 < A + B + C$

解法

　各コインを $a, b, c $ 枚持っている状態を考えます．この状態から，それぞれのコインが選ばれる（＝ $1$ 枚増える）確率は，

• $\frac a { a + b + c }$ で金貨
• $\frac b { a + b + c }$ で銀貨
• $\frac c { a + b + c }$ で銅貨

です．
　いずれかの種類のコインが $100$ 枚ある状態を「終端状態」と呼ぶことにします．「基本は全探索」という気持ちに立ち返ると，初期状態から始めてコインの増え方のパターンを全通り試したい気持ちがありますが，それは一般に指数通りあるのでできません．（言い方を変えれば，初期状態から終端状態へ到達する方法が沢山ある*1．）
　そこで，次のように状態をとる動的計画法を考え，同一視できる状態を潰します．

• $\mathit{ dp }_{ i, j, k } := \text{金貨が $i$ 枚，銀貨が $j$ 枚，銅貨が $k$ 枚ある状態に到達する確率}$

初期化は

• $\mathit{ dp }_{ A, B, C } = 1$

でそれ以外は $0$ です．遷移は，各 $\mathit{ dp }_{ i, j, k }$ から，先程の遷移確率を乗じつつ対応する行き先への足し込みになります．
　DP を回し終わったあと，ループを回して $i, j, k$ のいずれかが $100$ になっている状態へ至る操作回数の期待値を足し合わせます．期待値というのは対応する値とその値の発生確率の積です．ここで「値」というのは操作回数で，それは $( i - A ) + ( j - B ) + ( k - C )$ で求められます（コインが増えた回数の総和なので）．
　DP の遷移が $O( 1 )$ 個，コインの枚数が固定値なので状態数も $O( 1 )$ 個となり，$O( 1 )$ 時間です．

コード

int main()
{
	IN( int, A, B, C );
	static double dp[ 128 ][ 128 ][ 128 ];
	dp[A][B][C] = 1;

	REP( i, 100 )
	{
		REP( j, 100 )
		{
			REP( k, 100 )
			{
				if ( i == 0 && j == 0 && k == 0 )
				{
					continue;
				}

				{
					const double p = 1. * i / ( i + j + k );
					dp[ i + 1 ][j][k] += dp[i][j][k] * p;
				}
				{
					const double p = 1. * j / ( i + j + k );
					dp[i][ j + 1 ][k] += dp[i][j][k] * p;
				}
				{
					const double p = 1. * k / ( i + j + k );
					dp[i][j][ k + 1 ] += dp[i][j][k] * p;
				}
			}
		}
	}

	double res = 0;
	REP( i, 101 )
	{
		REP( j, 101 )
		{
			REP( k, 101 )
			{
				if ( i == 100 || j == 100 || k == 100 )
				{
					const int d = ( i - A )	+ ( j - B ) + ( k - C );
					res += d * dp[i][j][k];
				}
			}
		}
	}
	cout << res << endl;

	return 0;
}

*1:あんまり変わってないね