メモ帳がわり

個人的なメモを残します。主に競プロ

【ABC184】D - increment of coins

AtCoder AtCoder Beginner Contest DP 確率DP 期待値 確率 動的計画法

atcoder.jp

解法

再帰によって条件を満たすそれぞれの確率を求めようとすると、3100以上の計算量が必要になる場合がある。
愚直に解こうとすると指数時間になってしまう場合、DP(メモ化再帰)が有効になる場合は良くある。
そこで、3つの硬貨の数の取りうる場合の数は、最大で1003の通りしかないことに着目して、次のようなDPを考える。

dp[i][j][k] ← 金貨i枚, 銀貨j枚, 銅貨k枚の状態になる確率

初期状態は、 dp[A][B][C] ← 1

i, j, k枚の状態から、i+1, j, k枚に遷移する確率は、i/(i+j+k)なので、dp[i][j][k]にこの確率をかけると遷移できる。
また、i+1,j,k枚になる場合は1通りとは限らないことに注意すると漸化式は下のようになる。

dp[i+1][j][k] ← dp[i+1][j][k] + i/(i+j+k)
dp[i][j+1][k] ← dp[i][j+1][k] + i/(i+j+k)
dp[i][j][k+1] ← dp[i][j][k+1] + i/(i+j+k)

DPテーブルが完成したら期待値を求める。
硬貨が1つでも100枚になっている場合の確率だけを見ればよい。
硬貨が1つでも100枚になっている場合の数は、3 * 1002なので、2重ループで実装できそう。
一つの硬貨が100枚でその他の硬貨がそれぞれi, j枚の時の確率は下のように計算できる。
dp[100][i][j] + dp[i][100][j] + dp[i][j][100]
期待値を求めるには操作回数が必要なのでそれを求めると、100+i+j-(A+B+C)になる。

実装

void solve(long long A, long long B, long long C){
    vector<vector<vector<double>>> dp(102, vector<vector<double>>(102, vector<double>(102, 0)));
    dp[A][B][C] = 1;

    for (int i = A; i < 100; i++) {
        for (int j = B; j < 100; j++) {
            for (int k = C; k < 100; k++) {
                double I = i, J = j, K = k, N = i+j+k;
                dp[i + 1][j][k] += dp[i][j][k] * I/N;
                dp[i][j + 1][k] += dp[i][j][k] * J/N;
                dp[i][j][k + 1] += dp[i][j][k] * K/N;
            }
        }
    }



    double ans = 0;
    REP (i, 100) {
        REP (j, 100) {
            int n = 100+i+j-(A+B+C);
            ans += n * (dp[100][i][j] + dp[i][100][j] + dp[i][j][100]);
        }
    }

    cf(ans)

}

詰まったところ

i+1, j, kに遷移するには、i, j, kから遷移するしかないと思って下のように実装してしまっていた。
実際には、i+1, j-1, kや i+1, j, k-1から遷移することもできたのでバグを引き起こしてしまった。

for (int i = A; i <= 100; i++) {
    for (int j = B; j <= 100; j++) {
        for (int k = C; k <= 100; k++) {
            double I = i, J = j, K = k, N = i+j+k;
            dp[i + 1][j][k] = dp[i][j][k] * I/N;
            dp[i][j + 1][k] = dp[i][j][k] * J/N;
            dp[i][j][k + 1] = dp[i][j][k] * K/N;
        }
    }
}

上の実装を修正しても結果下のようになったがまだバグが潜んでいた。
for文の終了条件を100より大きくなること、としているが、これだとi=98, j = 100, k = 99の時などにおかしくなってしまう。
というのも、本来はどれかの硬貨が100枚になった時点で操作を終了しなくてはならないのだが、これだとその後も操作を続けることが可能になってしまう。
それでありえない遷移が起こってしまい、バグの原因になっていた。

for (int i = A; i <= 100; i++) {
    for (int j = B; j <= 100; j++) {
        for (int k = C; k <= 100; k++) {
            double I = i, J = j, K = k, N = i+j+k;
            dp[i + 1][j][k] += dp[i][j][k] * I/N;
            dp[i][j + 1][k] += dp[i][j][k] * J/N;
            dp[i][j][k + 1] += dp[i][j][k] * K/N;
        }
    }
}