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

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

I - Coins / Educational DP Contest

EDPC AtCoder C++ Educational DP Contest DP 確率 競技プログラミング

問題

n を正の奇数とする。

n 枚のコインがあり、コインには番号 1, \, 2, \, \cdots , \, n が振られている。 コイン i を投げるとき、コイン i は確率 {p}_{i} で表になり、確率 1 - {p}_{i} で裏になる。

n 枚すべてのコインを投げるとき、表の個数が裏の個数を上回る確率を求めよ。

入力

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

その次の1行に、確率 {p}_{1}, \, {p}_{2}, \, \cdots , \, {p}_{n} が順番に与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  • n は奇数である。
  • 3 \leq n \leq 2999
  • {p}_{i} は実数であり、小数第2位まで与えられる。
  • 0 \lt {p}_{i} \lt 1

出力

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

解法

\mathrm{dp}_{i, \, j} を、「コイン 1, \, \cdots, \, i を投げたときに、表が j 枚である確率」として定義します。 また、便宜上 \mathrm{dp}_{0, \, 0} = 1 とします。 （コインを0枚投げたときには表となる枚数は0枚なので、これは必ず満たされると言えます。）

ここで、 i \geq 1 として、コイン 1, \, \cdots , \, i - 1 までが投げられているという状態を仮定します。 また、 \mathrm{dp} の値については、 \mathrm{dp}_{0, \, j}, \, \mathrm{dp}_{1, \, j}, \, \cdots , \, \mathrm{dp}_{i - 1, \, j} までが計算されているとします。

このとき、「 i - 1 枚目まで投げて j 枚が表となっている」状況において、コイン i を投げたときの状況は以下のように分岐できると考えられます。

  • コイン i が表となり、表のコインの枚数が j + 1 となる（確率 {p}_{i} で発生）。
  • コイン i が裏となり、表のコインの枚数が j となる（確率 1 - {p}_{i} で発生）。

この2つの分岐から、 \mathrm{dp}_{i - 1, \, j} の値を基に、以下のように \mathrm{dp} の値を更新できると言えます。

\begin{gather} \mathrm{dp}_{i, \, j + 1} \leftarrow \mathrm{dp}_{i, \, j + 1} + \mathrm{dp}_{i - 1, \, j} \cdot {p}_{i} \\ \mathrm{dp}_{i, \, j} \leftarrow \mathrm{dp}_{i, \, j} + \mathrm{dp}_{i - 1, \, j} \cdot (1 - {p}_{i}) \end{gather}

よって、この更新を i = 1, \, 2, \, \cdots , \, n の順に行い、その中で j = 0, \, 1, \, \cdots , \, i の順番で行っていくことにより、すべての \mathrm{dp}_{n, \, j} の値が求められると言えます。

最終的に答えとして必要なものは、「表の枚数が裏の枚数よりも多い確率」です。 表の枚数が j 枚のとき、裏の枚数は n - j であるので、 j \gt n - j を満たすすべての整数 j に対して \mathrm{dp}_{n, \, j} の値の総和がこの問題の答えとなります。

この方法では、 \mathrm{dp}_{i, \, j} に対して、 1 \leq 1 \leq n, \, 0 \leq j \leq i のそれぞれの場合について更新をしていることから、全体で O({n}^{2}) 程度の計算量で導出できると言えます。 n \leq 2999 であるので、これは実行時間制限に十分間に合います。

ソースコード

int n;
double p[3005], dp[3005][3005];

int main() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> p[i];
  }

  // dp の値の更新
  dp[0][0] = 1;
  for (int i = 1; i <= n; i++) {
    for (int j = 0; j <= i; j++) {
      dp[i][j] += dp[i - 1][j] * ((double)1 - p[i]); // コイン i が裏の場合の更新
      dp[i][j + 1] += dp[i - 1][j] * p[i]; // コイン i が表の場合の更新
    }
  }

  double ans = 0; // 最終的な答えとなる確率
  for (int j = 0; j < n - j; j++) {
    ans += dp[n][n - j]; // 表が n - j 枚、裏が j 枚のときの確率を足していく（ n - j > j と設定しているため、これでOK）
  }

  // 値の出力（小数点の桁数を指定することを忘れない）
  cout << fixed << setprecision(32) << ans << endl;

  return 0;
}

感想

H問題 に続いて、記事を書くことができました。

確率を求める問題は、DPと相性がよかったりするイメージが個人的にはあります。

  • 前回: H問題
  • 次回: （未定）