Sushi [Educational DP Contest / DP まとめコンテスト J]

https://atcoder.jp/contests/dp/tasks/dp_j

前提知識

• 期待値DP

解説

https://atcoder.jp/contests/dp/submissions/3963610

もし考え方違ってたら指摘ください…

期待値DPをする。
何番目の寿司かということは特に関係なく、残っている個数毎に集計して問題ない。
すると、制約もあって、以下のDPが考えつく。
dp[c1][c2][c3] := 1個残りがc1個, 2個残りがc2個, 3個残りがc3個である状態にするまでの操作回数の期待値
　
mathさんの期待値と条件付確率を読む。
これの条件付き期待値を使って解いてみる。
　
dp[c1][c2][c3]について考える。
まず、確率pで次に遷移可能なときの期待値は1/pになる。
sm = c1+c2+c3とすると、sm/Nで状態が変わるので、N/smだけ遷移に必要。
その遷移先と確率は

• 1個の寿司を食べる dp[c1-1][c2][c3] 状態へ行く確率はc1/sm
• 2個の寿司を食べる dp[c1+1][c2-1][c3] 状態へ行く確率はc2/sm
• 3個の寿司を食べる dp[c1][c2+1][c3-1] 状態へ行く確率はc3/sm

となる。この遷移先が条件付き期待値になっているので、
dp[c1][c2][c3] = N/sm + dp[c1-1][c2][c3] × c1/sm + dp[c1+1][c2-1][c3] × c2/sm

1. dp[c1][c2+1][c3-1] × c3/sm

となる。
これを最後までやる。
c1=c2=c3=0のときはN/smで0割りになってしまうので、そのときだけcontinueを書いておこう。

int N, A[303], C[4];
double dp[303][303][303];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> A[i];
    rep(i, 0, N) C[A[i]]++;
 
    rep(c3, 0, N + 1) rep(c2, 0, N + 1) rep(c1, 0, N + 1) {
        int sm = c1 + c2 + c3;
        if (sm == 0) continue;
 
        dp[c1][c2][c3] = 1.0 * N / sm;
        if (c1) dp[c1][c2][c3] += dp[c1 - 1][c2][c3] * c1 / sm;
        if (c2) dp[c1][c2][c3] += dp[c1 + 1][c2 - 1][c3] * c2 / sm;
        if (c3) dp[c1][c2][c3] += dp[c1][c2 + 1][c3 - 1] * c3 / sm;
    }
 
    printf("%.10f\n", dp[C[1]][C[2]][C[3]]);
}