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.

// Idea : Probability = valid cases/all cases

// Example: [2, 1, 1]
// for first ball, we can put both ball in first bin and 0 ball in second bin OR we can put 1 ball in first bin, second ball in 2nd bin OR we can put both in second bin

// for second ball, we can put ball in first bin and 0 ball in second bin, similarly, we can put 1 ball in second bin.

// same thing with the third ball

// Try all possible combinations recursively. And, in the end check, if there are equal number of distinct balls in both bins or not.

// Note: Overflow will not happen with long. Why? Because, all possible cases (i.e denominator) =. n C (n/2). (n = number of balls) => 48 C 24 < 10^18

class Solution {
    public double getProbability(int[] balls) {
        // Get all possible cases where we select same number of balls in both bins
        long all = allCases(balls, 0, 0, 0);
        // Get all possible cases where we select same number of balls in both bins + we
        // select same number of distinct balls
        long valid = casesWithEqualDistinctBalls(balls, 0, 0, 0, 0, 0);

        return ((1.0) * valid / all);
    }

    // first = number of balls in first bin
    // second = number of balls in second bin
    public long allCases(int[] balls, int pos, int first, int second) {
        if (pos == balls.length) {
            // for all cases, we just need to check if both bins have same number of balls
            // or not
            if (first == second)
                return 1;
            return 0;
        }
        // Case ->1
        // we put all balls from current pos in second bin
        long answer = allCases(balls, pos + 1, first, second + balls[pos]);
        // Case ->2
        // we put all balls from current pos in first bin
        answer += allCases(balls, pos + 1, first + balls[pos], second);
        // Case ->3
        // Distributing balls between the 2 containers from current pos
        for (int i = 1; i < balls[pos]; ++i) {
            answer += (allCases(balls, pos + 1, first + i, second + balls[pos] - i) * nCr(balls[pos], i));
            // we multiply with nCr(b[pos], i)) because, even though we choosing balls of
            // same color, but they distinct with each other as well.
        }
        return answer;

    }

    // disF = distinct balls in first bin
    // disS = distinct balls in second bin
    // first = number of balls in first bin
    // second = number of balls in second bin
    public long casesWithEqualDistinctBalls(int[] balls, int pos, int first, int second, int disF, int disS) {
        if (pos == balls.length) {
            // we need to check if both bins have same number of balls or not + number of
            // distinct balls in each bin should be equal
            if (first == second && disF == disS)
                return 1;
            return 0;
        }
        // Case -> 1
        // we put all balls in second bin
        long answer = casesWithEqualDistinctBalls(balls, pos + 1, first, second + balls[pos], disF, disS + 1);
        // Case -> 2
        // we put all balls in first bin
        answer += casesWithEqualDistinctBalls(balls, pos + 1, first + balls[pos], second, disF + 1, disS);
        // Case ->3
        for (int i = 1; i < balls[pos]; ++i) {
            answer += (casesWithEqualDistinctBalls(balls, pos + 1, first + i, second + balls[pos] - i, disF + 1,
                    disS + 1) * nCr(balls[pos], i));
            // we multiply with nCr(b[pos], i)) because, even though we choosing balls of
            // same color, but they distinct with each other as well.
        }
        return answer;

    }

    long nCr(long n, long r) {
        return fact(n) / (fact(r) * fact(n - r));
    }

    long fact(long n) {
        long res = 1;
        for (int i = 2; i <= n; i++)
            res = res * i;
        return res;
    }
}