两个盒子中球的颜色相同的概率

1467. 两个盒子中球的颜色数相同的概率 - 力扣（LeetCode）

---

题型分析

概率DP+组合数

题目描述&数据范围

数据保证：
$$
1\leq balls.size() \leq 8\\
1\leq balls[i]\leq6
$$

题解

设 $dp[c][i][j]$ 表示前 $c$ 个颜色，第 $1$ 个盒子放置 $i$ 个球，第 $1$ 个盒子比第二个多 $j$ 种颜色的总情况

设 $ball[c]$ 表示第 $c$ 个颜色的球的数量，第 $1$ 个盒子分配 $r$ 个球，当前球的总数为 $sum$

考虑到 $j$ 可以取负数，所以加上总颜色数

则转移为：
$$
dp[c][i][j]=\begin{cases}
dp[c][i][j-1]\cdot C_{sum-i}^{ball[c]}&r=0\\
dp[c][i-r][j]\cdot C_{i}^{r}\cdot C_{sum-i}^{ball[c]-r}&0<r<ball[c]\\
dp[c][i][j+1]\cdot C_{i}^{ball[c]}&r=c
\end{cases}
$$

代码

//此处的转移利用加法的形式会更好写
class Solution {
public:
    double C[100][100],fac[60];
    void init(){
        for(int i=0;i<100;++i){
            for(int j=0;j<=i;++j){
                if(!j)C[i][j]=1;
                else C[i][j]=C[i-1][j]+C[i-1][j-1];
            }
        }
    }
    double getProbability(vector<int>& balls) {
        init();
        int sum=0,k=balls.size(),n=0;
        for(auto c:balls)n+=c;
        fac[0]=1;
        for(int i=1;i<=n;++i)fac[i]=fac[i-1]*i;//j
        double tot=fac[n];//总排列数
        vector<vector<double>> dp(n+1,vector<double>(2*k+1,0.0)); dp[0][k] = 1.0;
        for(auto c:balls){
            sum+=c;
            tot/=fac[c];
            vector<vector<double>> tmp(n+1,vector<double>(2*k+1,0.0));
            for(int r=0;r<=c;++r){
                int t=0;
                if(!r)t=-1;
                if(r==c)t=1;
                for(int i=0;i<=n;++i){
                    for(int j=0;j<=2*k;++j){
                        if(!dp[i][j])continue;//无贡献
                        double w=dp[i][j];
                        w*=C[i+r][r];
                        w*=C[sum-r-i][c-r];
                        tmp[i+r][j+t]+=w;
                    }
                }
            }
            swap(tmp,dp);
        }
        return dp[n/2][k]/tot;
    }
};