1467.Probability of a Two Boxes Having The Same Number of Distinct Balls

Tags: Math Backtracking Hard

Links: https://leetcode.com/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls/

---

Given 2n balls of k distinct colors. You will be given an integer array balls of size k where balls[i] is the number of balls of color i.

All the balls will be shuffled uniformly at random, then we will distribute the first n balls to the first box and the remaining n balls to the other box (Please read the explanation of the second example carefully).

Please note that the two boxes are considered different. For example, if we have two balls of colors a and b, and two boxes [] and (), then the distribution [a] (b) is considered different than the distribution [b] (a)(Please read the explanation of the first example carefully).

We want to calculate the probability that the two boxes have the same number of distinct balls.

Example 1:

Input: balls = [1,1]
Output: 1.00000
Explanation: Only 2 ways to divide the balls equally:
- A ball of color 1 to box 1 and a ball of color 2 to box 2
- A ball of color 2 to box 1 and a ball of color 1 to box 2
In both ways, the number of distinct colors in each box is equal. The probability is 2/2 = 1

Example 2:

Input: balls = [2,1,1]
Output: 0.66667
Explanation: We have the set of balls [1, 1, 2, 3]
This set of balls will be shuffled randomly and we may have one of the 12 distinct shuffles with equale probability (i.e. 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]
After that we add the first two balls to the first box and the second two balls to the second box.
We can see that 8 of these 12 possible random distributions have the same number of distinct colors of balls in each box.
Probability is 8/12 = 0.66667

Example 3:

Input: balls = [1,2,1,2]
Output: 0.60000
Explanation: The set of balls is [1, 2, 2, 3, 4, 4]. It is hard to display all the 180 possible random shuffles of this set but it is easy to check that 108 of them will have the same number of distinct colors in each box.
Probability = 108 / 180 = 0.6

Example 4:

Input: balls = [3,2,1]
Output: 0.30000
Explanation: The set of balls is [1, 1, 1, 2, 2, 3]. It is hard to display all the 60 possible random shuffles of this set but it is easy to check that 18 of them will have the same number of distinct colors in each box.
Probability = 18 / 60 = 0.3

Example 5:

Input: balls = [6,6,6,6,6,6]
Output: 0.90327

Constraints:

• 1 <= balls.length <= 8
• 1 <= balls[i] <= 6
• sum(balls) is even.
• Answers within 10^-5 of the actual value will be accepted as correct.

---

思路是使用古典概型，计算出所有的排列可能，然后用符合要求的种数除以总的排列种数。

两个盒子不同，模型可以等价于将2n个球排成一排，所有可能的排列就是 $$ C_n^{a_1} \times C_{n - a_1}^{a_2} \times C_{n - a_1 - a_2} ^{a_3} \cdots C_{a_k} ^{a_k} $$ 也就是先从所有的球里面，第一种颜色的球先选定位置，然后在剩下空出来的位置里选第二种颜色的球的位置。接下来计算所有满足要求的种数。那么可以考虑把每种颜色的球分成两组，其中一组可以是0，这样的划分我们就会得到两个数组firstBox和secondBox，长度和balls等长，第i位存储对balls中第i种颜色的球的划分，满足firstBox[i] + secondBox[i] = balls[i]。划分这种操作使用DFS再合适不过了，使用DFS需要考虑的两个问题：

• 什么样的划分是满足要求的？于是我们用leftSum和rightSum记录firstBox和secondBox内数值的总和，其中任意一个大于球总和的一半，后面就无需划分了，肯定不符合要求。
• 满足什么条件下输出结果？用pos记录划分到balls中的位置，当pos == balls.size()，意味着balls中所有的球都完成了划分，但是还需要满足一个条件，不同颜色的球的颜色数相同，意味着firstBox和secondBox内数据的非零个数相同。

因为要计算组合数，所以立刻会想到Pascal三角形，所以用calculate()函数来进行预处理，这样后面使用组合数的时候就可以直接得到了。如果对此有疑问，可以通过LeetCode 118体会。因为题目限定数组长度为8，每个数组最大值为6，所以最多需要计算48行，但是$C_{48}^6$是一个很大的数字，如果序列长度是8，数值全是6，那么即使用long long也会溢出，所以可以考虑用double来记录。

最开始一直最后一个测试用例过不去，然后发现题意理解错了，错误的以为“颜色数”相同不仅要颜色的种类数相同，还要对应数量相同。如果不需要对应数量相同，那么只需要用一个计数器记录非零元素的个数即可，无需排序操作了。

class Solution {
    int totalNum, halfNum;
    vector<int> firstbox, secondBox;
    vector<vector<double>> pascalTriangle;    

public:
    double getProbability(vector<int>& balls) {
        std::ios_base::sync_with_stdio(false);
        cin.tie(NULL);
        cout.tie(NULL);

        totalNum = accumulate(balls.begin(), balls.end(), 0);
        if (totalNum & 1) return 0;
        halfNum = totalNum >> 1;

        int n = balls.size();
        firstbox.resize(n), secondBox.resize(n);
        calculate(); //计算Pascal三角形

        //cout << combination(balls) << endl;

        return DFS(balls, 0, 0, 0) / combination(balls);
    }

    void calculate()
    {
        pascalTriangle.resize(49, vector<double>(49, 1));
        for (int i = 0; i <= 48; ++i) {
            for (int j = 1; j < i; ++j) {
                pascalTriangle[i][j] = pascalTriangle[i - 1][j - 1] + pascalTriangle[i - 1][j];
            }
        }
    }

    double combination(const vector<int> & nums)
    {
        double res = 1;
        int sum = accumulate(nums.begin(), nums.end(), 0);
        for (auto & e : nums) {
            if (e) {
                res *= pascalTriangle[sum][e];
                sum -= e;
            }
        }
        return res;
    }

    double DFS(const vector<int> & balls, int pos, int leftSum, int rightSum)
    {
        if (leftSum > halfNum || rightSum > halfNum) return 0;
        if (pos == (int)balls.size()) {
            // vector<int> tmpA = firstbox, tmpB = secondBox;
            // sort(tmpA.begin(), tmpA.end());
            // sort(tmpB.begin(), tmpB.end());

            //if (tmpA != tmpB) return 0;
            int cntA = 0, cntB = 0;
            for (auto & e : firstbox) if(e) ++cntA;
            for (auto & e : secondBox) if (e) ++cntB;
            if (cntA != cntB) return 0;

            return combination(firstbox) * combination(secondBox);
        }

        double res = 0;
        for (int i = 0; i <= balls[pos]; ++i) {
            firstbox[pos] = i;
            secondBox[pos] = balls[pos] - i;
            res += DFS(balls, pos + 1, leftSum + firstbox[pos], rightSum + secondBox[pos]);
        }

        return res;
    } 
};