Hard

题目描述

给你一个整数数组 balls，其中 balls[i] 是颜色为 i 的球的数量。

现在，你要将这些球放入两个盒子中，使得：

  • 每个盒子中球的总数相同
  • 求两个盒子中球的颜色数相同的概率

所有的球会被随机洗牌，然后前 n 个球放入第一个盒子，剩下的 n 个球放入第二个盒子。

请注意，两个盒子被认为是不同的。例如，如果我们有颜色为 a 和 b 的两个球，以及两个盒子 [] 和 ()，那么分布 [a] (b) 与分布 [b] (a) 被认为是不同的。

返回两个盒子中球的颜色数相同的概率。在实际值的 10^-5 以内的答案将被接受。

示例 1：

输入：balls = [1,1]
输出：1.00000
解释：只有 2 种方法来平均分配球：
- 颜色 1 的球放到盒子 1，颜色 2 的球放到盒子 2
- 颜色 2 的球放到盒子 1，颜色 1 的球放到盒子 2
在两种方法中，每个盒子中不同颜色的数量都相等。概率是 2/2 = 1

示例 2：

输入：balls = [2,1,1]
输出：0.66667
解释：我们有球的集合 [1, 1, 2, 3]

示例 3：

输入：balls = [1,2,1,2]
输出：0.60000

约束：

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

解题思路

这是一个概率统计问题，需要使用回溯和组合数学来解决。

核心思路是通过回溯枚举所有可能的分配方案，然后计算满足条件的概率。

具体步骤：

  1. 回溯枚举：对于每种颜色的球，枚举分配到第一个盒子的数量（0到该颜色球的总数）
  2. 约束条件：确保两个盒子中的球数相等（各为总数的一半）
  3. 概率计算：使用多项式定理计算每种分配方案的概率
  4. 统计结果：累计所有满足"两个盒子颜色数相同"条件的方案概率

关键点：

  • 使用组合数 C(n,k) = n!/(k!(n-k)!) 计算分配概率
  • 每种分配方案的概率 = (各颜色组合数乘积) / C(总数, n)
  • 需要预计算阶乘来快速计算组合数
  • 回溯时需要剪枝：如果当前盒子球数已超过一半，直接返回

时间复杂度主要取决于回溯的分支数，在给定约束下是可行的。

代码实现

class Solution {
private:
    vector<long long> factorial;
    
    void calcFactorial(int n) {
        factorial.resize(n + 1);
        factorial[0] = 1;
        for (int i = 1; i <= n; i++) {
            factorial[i] = factorial[i - 1] * i;
        }
    }
    
    long long combination(int n, int k) {
        if (k > n || k < 0) return 0;
        return factorial[n] / (factorial[k] * factorial[n - k]);
    }
    
    pair<double, double> dfs(vector<int>& balls, int index, vector<int>& box1, vector<int>& box2, int sum1, int sum2, int target) {
        if (sum1 > target || sum2 > target) return {0, 0};
        
        if (index == balls.size()) {
            if (sum1 != target || sum2 != target) return {0, 0};
            
            // 计算当前分配的概率
            long long ways = 1;
            int colors1 = 0, colors2 = 0;
            
            for (int i = 0; i < balls.size(); i++) {
                ways *= combination(balls[i], box1[i]);
                if (box1[i] > 0) colors1++;
                if (box2[i] > 0) colors2++;
            }
            
            double prob = (double)ways;
            return {prob, colors1 == colors2 ? prob : 0};
        }
        
        double totalWays = 0, validWays = 0;
        
        for (int i = 0; i <= balls[index]; i++) {
            box1[index] = i;
            box2[index] = balls[index] - i;
            
            auto result = dfs(balls, index + 1, box1, box2, sum1 + i, sum2 + balls[index] - i, target);
            totalWays += result.first;
            validWays += result.second;
        }
        
        return {totalWays, validWays};
    }
    
public:
    double getProbability(vector<int>& balls) {
        int totalBalls = 0;
        for (int ball : balls) {
            totalBalls += ball;
        }
        
        calcFactorial(totalBalls);
        
        vector<int> box1(balls.size()), box2(balls.size());
        auto result = dfs(balls, 0, box1, box2, 0, 0, totalBalls / 2);
        
        return result.second / result.first;
    }
};

class Solution:
    def getProbability(self, balls: List[int]) -> float:
        from math import factorial
        
        def combination(n, k):
            if k > n or k < 0:
                return 0
            return factorial(n) // (factorial(k) * factorial(n - k))
        
        def dfs(index, box1, box2, sum1, sum2, target):
            if sum1 > target or sum2 > target:
                return 0, 0
            
            if index == len(balls):
                if sum1 != target or sum2 != target:
                    return 0, 0
                
                # 计算当前分配的概率
                ways = 1
                colors1 = colors2 = 0
                
                for i in range(len(balls)):
                    ways *= combination(balls[i], box1[i])
                    if box1[i] > 0:
                        colors1 += 1
                    if box2[i] > 0:
                        colors2 += 1
                
                return ways, ways if colors1 == colors2 else 0
            
            total_ways = valid_ways = 0
            
            for i in range(balls[index] + 1):
                box1[index] = i
                box2[index] = balls[index] - i
                
                total, valid = dfs(index + 1, box1, box2, sum1 + i, sum2 + balls[index] - i, target)
                total_ways += total
                valid_ways += valid
            
            return total_ways, valid_ways
        
        total_balls = sum(balls)
        target = total_balls // 2
        
        box1 = [0] * len(balls)
        box2 = [0] * len(balls)
        
        total_ways, valid_ways = dfs(0, box1, box2, 0, 0, target)
        
        return valid_ways / total_ways

public class Solution {
    private long[] factorial;
    
    private void CalcFactorial(int n) {
        factorial = new long[n + 1];
        factorial[0] = 1;
        for (int i = 1; i <= n; i++) {
            factorial[i] = factorial[i - 1] * i;
        }
    }
    
    private long Combination(int n, int k) {
        if (k > n || k < 0) return 0;
        return factorial[n] / (factorial[k] * factorial[n - k]);
    }
    
    private (double, double) DFS(int[] balls, int index, int[] box1, int[] box2, int sum1, int sum2, int target) {
        if (sum1 > target || sum2 > target) return (0, 0);
        
        if (index == balls.Length) {
            if (sum1 != target || sum2 != target) return (0, 0);
            
            long ways = 1;
            int colors1 = 0, colors2 = 0;
            
            for (int i = 0; i < balls.Length; i++) {
                ways *= Combination(balls[i], box1[i]);
                if (box1[i] > 0) colors1++;
                if (box2[i] > 0) colors2++;
            }
            
            double prob = (double)ways;
            return (prob, colors1 == colors2 ? prob : 0);
        }
        
        double totalWays = 0, validWays = 0;
        
        for (int i = 0; i <= balls[index]; i++) {
            box1[index] = i;
            box2[index] = balls[index] - i;
            
            var result = DFS(balls, index + 1, box1, box2, sum1 + i, sum2 + balls[index] - i, target);
            totalWays += result.Item1;
            validWays += result.Item2;
        }
        
        return (totalWays, validWays);
    }
    
    public double GetProbability(int[] balls) {
        int totalBalls = balls.Sum();
        CalcFactorial(totalBalls);
        
        int[] box1 = new int[balls.Length];
        int[] box2 = new int[balls.Length];
        
        var result = DFS(balls, 0, box1, box2, 0, 0, totalBalls / 2);
        
        return result.Item2 / result.Item1;
    }
}

var getProbability = function(balls) {
    const n = balls.reduce((sum, count) => sum + count, 0) / 2;
    
    function factorial(num) {
        let result = 1;
        for (let i = 2; i <= num; i++) {
            result *= i;
        }
        return result;
    }
    
    function multinomial(counts) {
        const total = counts.reduce((sum, count) => sum + count, 0);
        let result = factorial(total);
        for (const count of counts) {
            result /= factorial(count);
        }
        return result;
    }
    
    let validCount = 0;
    let totalCount = 0;
    
    function backtrack(index, box1, box2, count1, count2) {
        if (count1 > n || count2 > n) return;
        
        if (index === balls.length) {
            if (count1 === n && count2 === n) {
                const distinctColors1 = box1.filter(x => x > 0).length;
                const distinctColors2 = box2.filter(x => x > 0).length;
                
                const ways = multinomial(box1) * multinomial(box2);
                totalCount += ways;
                
                if (distinctColors1 === distinctColors2) {
                    validCount += ways;
                }
            }
            return;
        }
        
        for (let i = 0; i <= balls[index]; i++) {
            box1[index] = i;
            box2[index] = balls[index] - i;
            backtrack(index + 1, box1, box2, count1 + i, count2 + balls[index] - i);
        }
    }
    
    backtrack(0, new Array(balls.length).fill(0), new Array(balls.length).fill(0), 0, 0);
    
    return validCount / totalCount;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型分析
时间复杂度O(∏(balls[i] + 1))，回溯遍历所有可能的分配方案
空间复杂度O(k + 总球数)，k为颜色数量，用于存储分配状态和阶乘数组