题目描述

桌面上有 2n 个颜色不完全相同的球，球上的颜色共有 k 种。给定一个大小为 k 的整数数组 balls，其中 balls[i] 是颜色为 i 的球的数量。

所有的球都已经 随机打乱顺序，前 n 个球放入第一个盒子，后 n 个球放入另一个盒子（请认真阅读样例 2 的解释部分）。

注意到，这两个盒子是不同的。例如，两个球颜色分别为 a 和 b，盒子分别为 [] 和 ()，那么 [a] (b) 和 [b] (a) 这两种分配方式是不同的（请认真阅读样例 1 的解释部分）。

请计算 两个盒子中球的颜色数相同 的情况的概率。

样例

输入：balls = [1,1]
输出：1.00000
解释：球平均分配的方式只有两种：
- 颜色为 1 的球放入第一个盒子，颜色为 2 的球放入第二个盒子
- 颜色为 2 的球放入第一个盒子，颜色为 1 的球放入第二个盒子
这两种分配，两个盒子中球的颜色数都相同。所以概率为 2/2 = 1。

输入：balls = [2,1,1]
输出：0.66667
解释：球的列表为 [1, 1, 2, 3]
随机打乱，得到 12 种等概率的不同打乱方案，每种方案概率为 1/12：
[1,1 / 2,3], [1,1 / 3,2], [1,2 / 1,3], [1,2 / 3,1],
[1,3 / 1,2], [1,3 / 2,1], [2,1 / 1,3], [2,1 / 3,1],
[2,3 / 1,1], [3,1 / 1,2], [3,1 / 2,1], [3,2 / 1,1]
然后，我们将前两个球放入第一个盒子，后两个球放入第二个盒子。
这 12 种可能的随机打乱方式中的 8 种满足 两个盒子中球的颜色数相同。
概率 = 8/12 = 0.66667

输入：balls = [1,2,1,2]
输出：0.60000
解释：球的列表为 [1, 2, 2, 3, 4, 4]。
要想显示所有 180 种随机打乱方案是很难的，
但只检查 两个盒子中球的颜色数相同 的 108 种情况是比较容易的。
概率 = 108 / 180 = 0.6 。

输入：balls = [3,2,1]
输出：0.30000
解释：球的列表为 [1, 1, 1, 2, 2, 3]。
要想显示所有 60 种随机打乱方案是很难的，
但只检查 两个盒子中球的颜色数相同 的 18 种情况是比较容易的。
概率 = 18 / 60 = 0.3 。

输入：balls = [6,6,6,6,6,6]
输出：0.90327

限制

• 1 <= balls.length <= 8
• 1 <= balls[i] <= 6
• sum(balls) 是偶数。
• 答案与真实值误差在 10^-5 以内，则被视为正确答案。

算法

(动态规划，组合数学) $O(m^2 + n^2m)$

1. 假设球的种类数为 $n$，球的总个数为 $m$。设状态 $f(i, j, k)$ 表示考虑了前 $i$ 种颜色的球，第一个盒子放了 $j$ 个球，第一个盒子与第二个盒子颜色的差值为 $k$ 时的方案数。球种类的有效下标从 1 开始，$j$ 的范围为 $[0, m / 2]$，$k$ 的范围为 $[-n, n]$。
2. 初始时，$f(0, 0, 0) = 1$，其余为 0。
3. 转移时，我们采用计算后继的方式，即状态 $f(i, j, k)$ 可以转移到 $i+1$ 层的状态。设 $tot$ 为前 $i$ 个球一共有多少个，假设第 $i+1$ 种颜色有 $c$ 个球。
   • 如果全部 $c$ 个球放入第二个盒子，则可以转移到 $f(i + 1, j, k - 1)$，由于第二个盒子目前有 $m/2 - (tot - j)$ 个空位置，需要放入 $c$ 个球，则有 $\binom{m / 2 - (tot - j)}{c}$ 种方法放入。
   • 如果有 $0 < l < c$ 个球放入第一个盒子，$c - l$ 个球放入第二个盒子，则可以转移到 $f(i + 1, j + l, k)$，根据乘法原理，有 $\binom{m / 2 - j}{l} \cdot \binom{m / 2 - (tot - j)}{c - l}$ 种方法放入。
   • 如果全部 $c$ 个球放入第一个盒子，则可以转移到 $f(i + 1, j, k + 1)$，同理有 $\binom{m / 2 - j}{c}$ 种方法放入。
4. 最终的方案数为 $f(n, m / 2, 0)$，但由于方案数过大，我们需要在过程中将其表示为概率。在每一步，我们有 $c$ 个颜色为 $i + 1$ 的球，目前已经安排了 $tot$ 个球，还有 $m-tot$ 个位置，将这 $c$ 个球放入 $m - tot$ 个位置总共有 $\binom{m - tot}{j}$ 种方案，所以在转移过程中，除以这个方案数，即可得到每一步的概率。

时间复杂度

• 预处理组合数需要 $O(m^2)$ 的时间。
• 动态规划的状态数为 $O(n^2m)$，转移时间为常数。
• 故总时间复杂度为 $O(m^2 + n^2m)$。

空间复杂度

• 需要额外 $O(m^2)$ 的空间存储组合数，额外 $O(n^2m)$ 的空间存储状态。
• 故总空间复杂度为 $O(m^2 + n^2m)$。

C++ 代码

#define LL long long

class Solution {
public:
    double getProbability(vector<int>& balls) {
        int n = balls.size(), m = 0;
        for (int i = 0; i < n; i++)
            m += balls[i];

        int s = m / 2;
        vector<vector<vector<double>>> f(
            n + 1, vector<vector<double>>(
                s + 1, vector<double>(
                    2 * n + 1, 0)));

        vector<vector<LL>> C(m + 1, vector<LL>(m + 1, 0));

        C[0][0] = 1;
        for (int i = 1; i <= m; i++) {
            C[i][0] = 1;
            for (int j = 1; j <= i; j++)
                C[i][j] = C[i - 1][j - 1] + C[i - 1][j];
        }

        int tot = 0;
        f[0][0][0 + n] = 1;

        for (int i = 0; i < n; i++) {
            for (int j = max(0, tot - s); j <= min(tot, s); j++)
                for (int k = -n; k <= n; k++) {
                    int c = balls[i];
                    double d = 1.0 / C[m - tot][c];

                    if (k > -n / 2)
                        f[i + 1][j][k - 1 + n] += f[i][j][k + n] *
                            d * C[s - (tot - j)][c];

                    for (int l = 1; l < c; l++)
                        if (j + l <= s)
                            f[i + 1][j + l][k + n] += f[i][j][k + n] *
                                d * C[s - j][l] * C[s - (tot - j)][c - l];

                    if (j + c <= s && k < n / 2)
                        f[i + 1][j + c][k + 1 + n] += f[i][j][k + n] *
                            d * C[s - j][c];
                }

            tot += balls[i];
        }

        return f[n][s][0 + n];
    }
};