Leetcode刷题Python1467. 两个盒子中球的颜色数相同的概率

Posted 2022-10-14 Better Bench

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

1 题目

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

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

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

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

示例 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 。

示例 4：

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

示例 5：

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

提示：

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

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

2 解析

参考知乎文章讲解：https://zhuanlan.zhihu.com/p/145606485
困难题

3 Python实现

def get_num(self,balls):
    import math
    #计算重复排列数
    #重复排列数为：S!/(S1!*S2!*S3!*...)
    def get_num(self,balls):
        summ = math.factorial(sum(balls))
        temp = 1
        for ball in balls:
            temp *= math.factorial(ball) 
        return summ/temp
    #dfs遍历所有可能的分配，c为从第c颜色开始，balls为各个颜色的数量，left_arr存放该方案左边
    #部分的各个颜色数量，right_arr表示右边各个颜色的数量，left_sum和right_sum表示左边和右边
    #的球的数量，n为边界，两部分最多只能有n个球
    def dfs(self,c,balls,left_arr,right_arr,left_sum,right_sum,n):
        #如果某部分已经超过n，剪枝
        if left_sum > n or right_sum > n:
            return 0
        #分配结束，结束递归，看这种方案的重复排列数
        if c == len(balls):
            left_color = 0
            right_color = 0
            for x in left_arr:
                if x:
                    left_color += 1
            for x in right_arr:
                if x:
                    right_color += 1
            #统计左右部分颜色数是不是相等，相等才符合
            if left_color == right_color:
                #两边各满足重复排列数，满足的方案数为两边的重复排列数相乘
                return self.get_num(left_arr) * self.get_num(right_arr)
            else:
                #颜色数不相等，剪枝
                return 0

        res = 0
        #遍历这种颜色的球可能取值范围，i个分给左边，剩下的分给右边
        for i in range(balls[c]+1):
            left_arr[c] = i
            right_arr[c] = balls[c] - i
            res += self.dfs(c+1,balls,left_arr,right_arr,left_sum+left_arr[c],right_sum+right_arr[c],n)
        #res为所有分配方案的总和
        return res

    def getProbability(self, balls: List[int]) -> float:
        total = self.get_num(balls) #所有的重复排列数
        print(total)
        n = sum(balls) // 2 #n为边界，两边最多n个球
        #初始化左右数组，保存左边和右边的各颜色球个数
        left_arr = [0 for _ in range(len(balls))]
        right_arr = [0 for _ in range(len(balls))]
        #返回概率为符合条件方案数除以总的方案数
        return self.dfs(0,balls,left_arr,right_arr,0,0,n)/total