LeeCode 1467 dfs + 组合数

题意

传送门 LeeCode 1467. 两个盒子中球的颜色数相同的概率

题解

$dfs(i,lsum,lco)$ 为已处理前 $i-1$ 种颜色小球，当前第 $1$ 个盒子小球数量为 $lsum$，第 $1$ 个盒子与第 $2$ 个盒子颜色总数差值为 $lco$ 时，最终转移到合法状态的概率。

对于第 $i$ 种颜色的小球，枚举放 $0$ 到 $balls[i]$ 个小球到第一个盒子的情况，设此数量为 $n$，向下一种颜色小球转移的概率为 $C_{balls[i]}^n/2^{balls[i]}$，则

$$lc{o}_{i+1}=\{\begin{array}{ll}-1& n=0\\ 1& n=balls[i]\\ 0& otherwise\end{array}$$

$$dfs(i,lsum,lco)=\sum C_{balls[i]}^n/2^{balls[i]}dfs(i+1,lsum+n,lc{o}_{i+1})$$

$dfs$ 求得的结果是两个盒子小球数相同且球的颜色数相同的概率，设为 $P(x)$；而题意为求两个盒子小球数相同时，两个盒子球的颜色数相同的条件概率；设两个盒子小球数相同的概率为 $P(y)$，小球总数为 $sum$

$$P(y)=\frac{C_{sum}^{sum/2}}{2^{sum}}$$

则答案为

$$P(x|y)=\frac{P(xy)}{P(y)}=\frac{P(x)}{P(y)}$$

class Solution
{
public:
    int C[7][7];
    int fac[7], n, half;
    vector<int> balls;
    void init()
    {
        fac[0] = 1;
        for (int i = 1; i < 7; i++)
            fac[i] = 2 * fac[i - 1];
        for (int i = 0; i < 7; i++)
            C[i][0] = 1;
        for (int i = 1; i < 7; i++)
        {
            for (int j = 1; j < 7; j++)
            {
                C[i][j] = C[i - 1][j - 1] + C[i - 1][j];
            }
        }
    }
    double dfs(int b, int lnum, int lco)
    {
        if (lnum > half)
            return 0;
        if (b == n)
            return lnum == half ? (lco == 0 ? 1 : 0) : 0;
        double res = 0, m = fac[balls[b]];
        for (int i = 0; i <= balls[b]; i++)
        {
            if (i == 0)
            {
                res += 1.0 / m * dfs(b + 1, lnum, lco - 1);
            }
            else if (i == balls[b])
            {
                res += 1.0 / m * dfs(b + 1, lnum + i, lco + 1);
            }
            else
            {
                res += C[balls[b]][i] / m * dfs(b + 1, lnum + i, lco);
            }
        }
        return res;
    }
    double fact(int l, int r)
    {
        if (l == r)
            return l;
        return l * fact(l - 1, r);
    }
    double calc(int n)
    {
        double sum = 1;
        for (int i = 0; i < n; i++)
            sum *= 2;
        double c = fact(n, n / 2 + 1) / fact(n / 2, 1);
        return c / sum;
    }
    double getProbability(vector<int> &balls)
    {
        init();
        int n = balls.size(), sum = 0;
        for (int i = 0; i < n; i++)
            sum += balls[i];
        this->n = n, this->half = sum / 2, this->balls = balls;
        double p = calc(sum);
        return dfs(0, 0, 0) / p;
    }
};