yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

J - Sushi / Educational DP Contest

問題

 n 枚の皿があり、皿には  1, \, 2, \, \cdots , \, n と番号が振られている。

最初、各  i について、皿  i には  {a}_{i} 個の寿司が置かれている。

すべての寿司がなくなるまで、以下の操作を繰り返し行うとき、すべての寿司がなくなるまでの操作回数の期待値を求めよ。

  •  1, \, 2, \, \cdots , \, n の目が等確率で出るサイコロを振り、出目の番号の皿に置かれている寿司を1つ食べる。
    • ただし、その皿に寿司がない場合は何も行わない。

入力

まず最初の1行目に、整数  n が与えられる。

その次の1行に、整数  {a}_{1}, \, {a}_{2}, \, \cdots , \, {a}_{n} が順に与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  •  1 \leq n \leq 300
  •  1 \leq {a}_{i} \leq 3

出力

答えとなる値を1行で出力すること。 ただし、答えとの絶対誤差が  {10}^{-9} 以下ならば正解となる。

解法

皿はサイコロで等確率で選ばれるため、どの皿に何個の寿司が置かれているのか、ということは状態変化に関係がありません。

寿司が1個、2個、3個と置かれている皿の枚数をそれぞれ  i, \, j, \, k とすると、1個、2個、3個の寿司が置かれている皿が選ばれる確率はそれぞれ  \displaystyle \frac{i}{n}, \, \frac{j}{n}, \, \frac{k}{n} ということになります。 従って、寿司が置かれていない皿が選ばれる確率は  \displaystyle \left( 1 - \frac{i + j + k}{n} \right) となります。

このことから、寿司が1個、2個、3個と置かれている皿の枚数がそれぞれ  i, \, j, \, k のときに操作を1回行うとき、状態変化としては、

  • 確率  \displaystyle \frac{i}{n} で、寿司が1個、2個、3個と置かれている皿の枚数はそれぞれ  i - 1, \, j, \, k となる
  • 確率  \displaystyle \frac{j}{n} で、寿司が1個、2個、3個と置かれている皿の枚数はそれぞれ  i + 1, \, j - 1, \, k となる
  • 確率  \displaystyle \frac{k}{n} で、寿司が1個、2個、3個と置かれている皿の枚数はそれぞれ  i, \, j + 1, \, k - 1 となる
  • 確率  \displaystyle 1 - \frac{i + j + k}{n} で、寿司が1個、2個、3個と置かれている皿の枚数はそれぞれ  i, \, j, \, k となる

というものになると言えます。

ここで、  \mathrm{dp}_{i, \, j, \, k} を「寿司が1個、2個、3個と置かれている皿の枚数がそれぞれ  i, \, j, \, k のときの、最終的に寿司をすべてなくすまでの操作回数の期待値」として定めます。 このとき、以上の考察での変化後の状態では、最終的な操作回数が  1 回減ると考えられるので、 \begin{align} \mathrm{dp}_{i, \, j, \, k} = 1 + \left\{ \frac{i}{n} \cdot \mathrm{dp}_{i - 1, \, j, \, k} \right. + \frac{j}{n} \cdot \mathrm{dp}_{i + 1, \, j - 1, \, k} & + \frac{k}{n} \cdot \mathrm{dp}_{i, \, j + 1, \, k - 1} \\ & \left. + \left( 1 - \frac{i + j + k}{n} \right) \cdot \mathrm{dp}_{i, \, j, \, k} \right\} \end{align} という等式が成立すると言えます。 この等式を整理すると、 \begin{align} \frac{i + j + k}{n} \cdot \mathrm{dp}_{i, \, j, \, k} = 1 + \frac{i}{n} \cdot \mathrm{dp}_{i - 1, \, j, \, k} + \frac{j}{n} \cdot \mathrm{dp}_{i + 1, \, j - 1, \, k} + \frac{k}{n} \cdot \mathrm{dp}_{i, \, j + 1, \, k - 1} \end{align} より、 \begin{align} \mathrm{dp}_{i, \, j, \, k} = \frac{1}{i + j + k} \left( n + i \cdot \mathrm{dp}_{i - 1, \, j, \, k} + j \cdot \mathrm{dp}_{i + 1, \, j - 1, \, k} + k \cdot \mathrm{dp}_{i, \, j + 1, \, k - 1} \right) \end{align} が成立します。

また、どの皿にも寿司が置かれていない場合は操作を行う必要がないので、  \mathrm{dp}_{0, \, 0, \, 0} = 0 が成立します。

従って、初期状態での寿司が1個、2個、3個と置かれている皿の枚数をそれぞれ  i, \, j, \, k とするときに、以上の漸化式を用いて求められる  \mathrm{dp}_{i, \, j, \, k} の値が、この問題の答えとなります。

この問題での皿の枚数は  n で与えられるため、値の導出に必要な計算量は  O({n}^{3}) です。  n \leq 300 であるので、これは実行時間制限に十分間に合うと考えられます。

ソースコード

寿司が1個、2個、3個と置かれている皿の枚数をそれぞれ  i, \, j, \, k とするときの、  \mathrm{dp}_{i, \, j, \, k} の値を再帰的に求める関数 solve() を以下のように実装しました。

double dp[305][305][305];
bool visited[305][305][305]; // visited[i][j][k] は solve(i, j, k) を既に計算したかどうかを表す

double solve(int i, int j, int k) {
  // i, j, k が 0 未満のとき自体が発生しないので、このときは 0 を返す
  if (i < 0 || j < 0 || k < 0) {
    return 0;
  }

  // i = j = k = 0 のときは、操作回数の期待値は 0 なので、 0 を返す
  if (i == 0 && j == 0 && k == 0) {
    return 0;
  }
  // solve(i, j, k) を計算済であれば、そのまま dp[i][j][k] を返す
  if (visited[i][j][k]) {
    return dp[i][j][k];
  }

  // solve(i, j, k) を計算済ということにする
  visited[i][j][k] = true;

  // ans は最終的な期待値の値を表す
  double ans = (double)(n) / (double)(i + j + k);

  // 寿司が1個置かれている皿を選ぶ場合
  ans += (double)(i) / (double)(i + j + k) * solve(i - 1, j, k);

  // 寿司が2個置かれている皿を選ぶ場合
  ans += (double)(j) / (double)(i + j + k) * solve(i + 1, j - 1, k);

  // 寿司が3個置かれている皿を選ぶ場合
  ans += (double)(k) / (double)(i + j + k) * solve(i, j + 1, k - 1);

  // 最終的に dp[i][j][k] を ans で更新して、 ans を返す
  return dp[i][j][k] = ans;
}

感想

期待値のDPについては、状態変化を正確に把握しきれるかがポイントになると思います。 (これが結構難しいんですよね…。)

  • 前回: I問題
  • 次回: (未定)