AtCoder Educational DP Contest J問題解いてみた

今回の問題

問題文

$N$ 枚の皿があります。皿には $1, 2, \ldots, N$ と番号が振られています。
最初、各 $i$ ($1 \leq i \leq N$) について、皿 $i$ には $a_i$ ($1 \leq a_i \leq 3$) 個の寿司が置かれています。

すべての寿司が無くなるまで、太郎君は次の操作を繰り返し行います。

1. $1, 2, \ldots, N$ の目が等確率で出るサイコロを振り、出目を $i$ とする。
2. 皿 $i$ に寿司がある場合、皿 $i$ の寿司を $1$ 個食べる。
3. 皿 $i$ に寿司が無い場合、何も行わない。

すべての寿司が無くなるまでの操作回数の期待値を求めてください。

制約

• 入力はすべて整数である。
• $1 \leq N \leq 300$
• $1 \leq a_i \leq 3$

自分の回答

#include <bits/stdc++.h>
#include <atcoder/dsu>
using namespace std;
using namespace atcoder;
double dp[301][301][301];
int main() {
    int n;
    cin >> n;
    int alist[n];
    int blist[4] = {0,0,0,0};
    for (int i = 0; i < n; i++) {
        cin >> alist[i];
        blist[alist[i]]++;
    }
    
    memset(dp, 0, sizeof(dp));
    dp[0][0][0] = 0.0;
    for (int k = 0; k <= n; k++){
        for (int j = 0; j <= n; j++) {
            for (int i = 0; i <= n; i++) {
                if (i == 0 && j == 0 && k == 0) continue;
            
                if (i + j + k > n) continue;

                double res = (double)n; 
                if (k > 0) res += dp[k-1][j+1][i] * k;

                if (j > 0) res += dp[k][j-1][i+1] * j;

                
                if (i > 0) res += dp[k][j][i-1] * i;

                
                dp[k][j][i] = res / (i + j + k);
            }
        }
    }
    cout << fixed << setprecision(10) << dp[blist[3]][blist[2]][blist[1]] << endl;

    return 0;
}

解説

今回寿司の順番はどうでもよい。残りのすしが3皿、2皿、1皿の枚数がわかっていればいい。最初dp[0][0][0]は1である。ここから、もらうdpで、i,j,kの状態になるためには、どういう状態からなるかを考える。例えば321になるには411からや330からなどが考えられる。また、期待値の計算は一度紙で書いてから実装しないとわからない。