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

难度: Hard

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

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

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

请返回「两个盒子中球的颜色数相同」的情况的概率。答案与真实值误差在 10^-5 以内，则被视为正确答案

示例 1：

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

示例 2：

输入：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

示例 3：

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

提示：

1 <= balls.length <= 8
1 <= balls[i] <= 6
sum(balls) 是偶数

Submission

运行时间: 31 ms

内存: 16.2 MB

class Solution:
    def getProbability(self, balls: List[int]) -> float:
        def combine(x, y):#组合数
            return math.factorial(x) // math.factorial(y) // math.factorial(x - y)
        s = sum(balls)#球总数
        l = len(balls)#颜色总数
        @lru_cache(None)
        def dfs(i, c, t):
            if i == l:#如果颜色选完了
                return int(t == 0 and c == s // 2) #选了1/2的球数量且颜色无变化
            res = dfs(i + 1, c, t + 1) + dfs(i + 1, c + balls[i], t - 1)#不选和全选组合数都为1，直接相加即可，需要更新颜色变化
            for j in range(1, balls[i]):#其他情况，颜色无变化
                res += dfs(i + 1, c + j, t) * combine(balls[i], j)
            return res
        res = dfs(0, 0, 0)
        return res / combine(s, s // 2)

Explain

该题解采用了深度优先搜索（DFS）配合动态规划的方法来解决问题。首先，定义一个递归函数 `dfs(i, c, t)`，其中 `i` 表示当前处理到的球的颜色索引，`c` 表示当前已选取的球的颜色数量，`t` 表示当前两个盒子中球数量的差值。递归的基础情况是当所有颜色都处理完毕时，检查是否恰好选取了一半的球并且两个盒子中球的颜色数相同。对于每种颜色的球，可以选择不放任何球到第一个盒子、放全部球到第一个盒子，或者放部分球到第一个盒子。通过组合数计算这些选择的可能性，并递归地探索所有可能的分配方式。最后，计算满足条件的分配方式与总的分配可能性的比值，得到最终的概率。

时间复杂度: O(6^k)

空间复杂度: O(k * s)

class Solution:
    def getProbability(self, balls: List[int]) -> float:
        def combine(x, y): # 计算组合数 C(x, y)
            return math.factorial(x) // math.factorial(y) // math.factorial(x - y)
        s = sum(balls) # 球总数
        l = len(balls) # 颜色总数
        @lru_cache(None)
        def dfs(i, c, t):
            if i == l: # 所有颜色处理完毕
                return int(t == 0 and c == s // 2) # 检查是否选取了一半的球并且两盒颜色数相同
            res = dfs(i + 1, c, t + 1) + dfs(i + 1, c + balls[i], t - 1) # 不放球和放所有球的情况
            for j in range(1, balls[i]): # 放部分球的情况
                res += dfs(i + 1, c + j, t) * combine(balls[i], j)
            return res
        res = dfs(0, 0, 0)
        return res / combine(s, s // 2) # 计算概率

Explore

在递归函数 `dfs` 中，参数 `t` 的处理是关键以确保最终两个盒子中的球数相等。这个差值 `t` 是动态维护的，每次递归调用时，根据将球放入第一个盒子或第二个盒子中，`t` 会相应地增加或减少。例如，如果选择不放任何球到第一个盒子中，我们考虑 `t + 1`（因为第二个盒子相对多了一个球），如果放置所有球到第一个盒子，则考虑 `t - 1`（因为第一个盒子相对多了一个球）。最终，我们想要的是在递归结束时 `t` 等于0，即两个盒子中的球数完全相等。这是通过在每一步递归中正确调整 `t` 的值，并在递归的基础情况中检查 `t == 0` 来实现的。

在递归处理中，当选择放置部分球到第一个盒子时，这种情况并不会改变两个盒子之间球数量的总差 `t`，因为我们只是在第一个盒子中放入了部分球而没有改变总体的球数平衡。具体来说，`t + 1` 和 `t - 1` 是处理极端情况，即一个盒子比另一个盒子多一个球或者少一个球。而当放入部分球时，我们实际上是在两个盒子之间重新分配球的数量，但这种分配不影响两盒子间的总差异，因此在这种情况下 `t` 的值保持不变。

组合数计算 `combine(balls[i], j)` 是用来计算从 `balls[i]` 个球中选取 `j` 个球的方案数。这个计算在算法中非常关键，因为它决定了不同的球分配方案的数量。使用组合数允许我们精确地计算出每种分配可能性的数量，从而确保所有可能的情况都被考虑到。然而，组合数的计算是基于阶乘的，这可能在数值上非常大，尤其是当 `balls[i]` 的值较大时。虽然这增加了计算的复杂度，但通过使用缓存机制（如在解中使用的 `lru_cache` 或预计算组合数并存储）可以显著提高效率。如果未优化，组合数的频繁计算可能成为算法性能的瓶颈。