吃寿司，概率/期望 DP

[原题链接（https://atcoder.jp/contests/dp/tasks/dp_j）]

code

#include<bits/stdc++.h>
using namespace std;
//"O campeão tem nome, e se chama Charles Oliveira!"
#define int long long
#define endl '\n'
#define ep emplace
#define pob 
#define ll long long
#define pb push_back
#define pof pop_front
#define pob pop_back
#define all(a) a.begin(),a.end()
#define rall(a) a.rbegin(),a.rend()
#define mod 998244353
#define MOD 1000000007

using ld = long double;
using ui = unsigned;
using ull = unsigned long long;
using i128 = __int128;
const int MAXN=310;
//概率/期望dp
int n,a[5];
double dp[MAXN][MAXN][MAXN];
void solve()
{
	cin>>n;
	for(int i=1;i<=n;i++)
	{
		int tmp;
		cin>>tmp;
		a[tmp]++;
	}
    /*
    公式推导：

    这个1的意思是，不管你有没有吃，你都要就是加一步

    dp[i,j,k]=
    1+[n-(i+j+k)]/n*dp[i,j,k]
    +i/n*dp[i-1,j,k]
    +j/n*dp[i+1,j-1,k]
    +k/n*dp[i,j+1,k-1]

    dp不可以有依赖自己的，所以要dp[i,j,k]-[n-(i+j+k)]/n*dp[i,j,k]==(i+j+k)/n*dp[i,j,k]
    即s/n*dp[i,j,k],然后再把这个s/n再乘过去
    */

    //本题要求的就是 类似于加权平均，离散型随机变量的一切可能的取值  与对应的概率  乘积之和称为该离散型随机变量的数学期望 
    //就是1+未来所有可能局面的期望步数的加权平均值

    //dp 代表当桌子上有 𝑖个盘子剩 1 个寿司，𝑗个盘子剩 2 个寿司，𝑘个盘子剩 3 个寿司时，直到把桌上所有寿司全部吃完，期望还需要掷多少次骰子（操作多少步）。
	for(int k=0;k<=n;k++)
	{
		for(int j=0;j<=n;j++)
		{
			for(int i=0;i<=n;i++)
			{
            	double s=i+j+k;
				if(s!=0)//防止 i,j,k 同时为零,因为下面有i-1 ，j-1，k-1
				{
					if(i)dp[i][j][k]+=dp[i-1][j][k]*i/s;
        			if(j)dp[i][j][k]+=dp[i+1][j-1][k]*j/s;
    	    		if(k)dp[i][j][k]+=dp[i][j+1][k-1]*k/s;
	        		dp[i][j][k]+=(double)n/s;
         		}
			}
		}
	}
	cout<<fixed<<setprecision(10)<<dp[a[1]][a[2]][a[3]]<<endl;
}

signed main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    int t=1;
    //cin>>t;
    while(t--)solve();
}