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

Problem Explanation

In this problem, we have a specific number of balls of various colors. Our goal is to randomly partition these balls in two distinct boxes such that the number of distinct colored balls in both boxes are the same. We need to return the probability of this event.


To solve this problem, we will simulate the event using recursion. For every color, we are allowed to place 0 to balls[i] balls in the first box. Then we will calculate the number of balls we placed in the second box, which is balls[i] minus the balls we placed in the first box.

We assume we have two boxes A and B.

  1. Travel through the 'balls' array.

    • For every color, there are balls[i] + 1 types of splitting cases. We will use a loop to go through these cases from 0 to balls[i] and count the probability of each combination.
  2. Calculate the new number of balls for each box and also the color count.

    • If we took balls from a color, then the color count of the box increases by 1.
  3. Recurse to the next color and update counts.

    • ballsCountA, ballsCountB are the total number of balls in boxes A and B.

    • colorsCountA, colorsCountB are the total number of different colors in boxes A and B.

    • If ballsCountA or ballsCountB is greater than n, return 0, implying that either box has more balls than allowed.

  4. Check the return condition.

    • If we have checked all colors, then check the boxCase.

      • For boxCase kEqualBalls, we return 1 since here we consider the total number of color combinations, regardless of the number of distinct colors.

      • For boxCase kEqualDistantBalls, we need to check if the number of different colors in both boxes is the same. If yes, return 1; otherwise, return 0.

  5. Finally, return ans divided by (fact[ballsTakenA] * fact[ballsTakenB]), where 'fact' is an array of factorial values from 0! to 6!. This step is needed to avoid over-counting because the order of balls of the same color does not matter.

Python Solution

3from math import factorial as f
5class Solution:
6    def getProbability(self, balls):
7        self.n = sum(balls)
8        self.f = [f(i) for i in range(self.n+1)]
9        self.balls = balls
10        self.total = self.f[self.n] * self.dfs(self.n//2, self.n//2, 0, 0, 0)/self.combs(0, self.n)
11        return self.cases/self.total
13    def combs(self,i,n):
14        if i==len(self.balls):
15            return 1 if n==0 else 0
16        return sum(self.f[n-j]*self.combs(i+1, j) for j in range(min(n, self.balls[i])+1))
18    def dfs(self, a, b, ca, cb):
19        if ca-cb>1 or ca-cb<-1: return 0
20        if a==0 and b==0: return int(ca==cb)
21        if not self.balls: return 0
22        x, self.balls = self.balls[-1], self.balls[:-1]
23        return sum(f(a)*f(b)*self.dfs(a-na, b-nb, ca-int(na>0), cb-int(nb>0)) for na in range(min(a,x)+1) for nb in range(min(b, x-na)+1))

Here we use the math.factorial function to compute factorials. The getProbability function first computes the total number of combinations for the given set of balls using a helper function, dfs(), which uses Depth-First Search (DFS) on the list of balls recursively until an empty list is encountered. The values a and b represent the remaining number of balls that can be placed in boxes A and B respectively. The values ca and cb are the number of distinct colors in each box. We exit early if the difference in the number of colors in both boxes is more than 1, as we will never reach an equal amount beyond that point. If no balls remain, we check if both boxes have an equal color count and return the result as an integer.

Java Solution

3class Solution {
4    static final int MAX = 50;
5    int[] color;
6    double[][][][][] memo;
7    double all;
8    double f[] = new double[MAX + 1];
9    public double getProbability(int[] balls) {
10        f[0] = 1;
11        for (int i = 1; i <= MAX; i++) {
12            f[i] = f[i - 1] * i;
13        }
14        int sum = 0;
15        for (int b : balls) {
16            sum += b;
17        }
18        color = balls;
19        memo = new double[balls.length][sum / 2 + 1][balls.length + 1][sum / 2 + 1][balls.length + 1];
20        all = f[sum];
21        for (int b : balls) {
22            all /= f[b];
23        }
24        return divide(0, 0, 0, 0, 0, sum / 2);
25    }
27    double divide(int i, int s1, int c1, int s2, int c2, int half) {
28        if (i >= color.length) {
29            if (s1 != s2 || c1 != c2) return 0;
30            return 1;
31        }
32        if (memo[i][s1][c1][s2][c2] != 0) {
33            return memo[i][s1][c1][s2][c2] - 1;
34        }
35        for (int a = 0; a <= color[i]; a++) {
36            if (s1 + a <= half) {
37                int nc1 = c1 + (a > 0 ? 1 : 0);
38                for (int b = 0; b <= color[i]; b++) {
39                    if (a + b != color[i] || s2 + b > half) continue;
40                    int nc2 = c2 + (b > 0 ? 1 : 0);
41                    double way = f[a] * f[b] / f[color[i]];
42                    double ans = divide(i + 1, s1 + a, nc1, s2 + b, nc2, half) * way;
43                    memo[i][s1][c1][s2][c2] += ans;
44                }
45            }
46        }
47        return memo[i][s1][c1][s2][c2]++;
48    }

In this Java solution, we store the factorial values in an array f during initialization. We use memoization to store the computed probabilities. The function divide is recursively called to calculate the total number of combinations that meet our criteria (equal number of unique colors in both boxes). The variables s1, s2, c1, c2, represent the number of balls and distinct colors in each box. If we've processed all colors, we check if we have an equal number of unique colors and total balls in both boxes.

Javascript Solution

3var getProbability = function(balls) {
4  const sum = balls.reduce((a, b) => a + b, 0);
5  const comb = Array.from(Array(sum + 1), () => new Array(sum + 1).fill(0));
7  comb[0][0] = 1;
8  for (let s = 1; s <= sum; s++) {
9      comb[s][0] = 1;
10      for (let i = 1; i <= s; i++) {
11          comb[s][i] = comb[s-1][i-1] + comb[s-1][i];
12      }
13  }
15  const dfs = (balls, a, b, ca, cb) => {
16      if (!balls.length) {
17          return +(a === b && ca === cb);
18      }
20      const x = balls[0];
21      let res = 0;
22      for (let A = 0; A <= Math.min(a, x); A++) {
23          for (let B = 0; B <= Math.min(b, x - A); B++) {
24              res += comb[x][A] * comb[x - A][B] * comb[a][A] * comb[b][B] * dfs(balls.slice(1), a - A, b - B, ca + (A > 0), cb + (B > 0));
25          }
26      }
28      return res;
29  };
31  return dfs(balls, sum / 2, sum / 2, 0, 0) / comb[sum][sum / 2];

In this Javascript solution, we compute the number of combinations for each value from 0 to sum. We then recursively traverse through the balls array, for each ball, we calculate the number of ways to distribute the balls into the two boxes while keeping track of the number of balls and distinct ball colors in each box. We ensure that no box gets more than half of the total number of balls, and we also check if the boxes have the same number of distinct colors after distributing all the balls before updating the result. Finally, we return the result divided by the total number of combinations.The Javascript solution uses the combinatorial mathematical principle to create a two-dimensional array to store all potential combinations of balls. The dfs (depth-first search) function is defined to handle the recursive logic of traversing through all possible distributions of balls in the boxes. It takes balls, an array of integers where balls[i] is the number of balls of the i-th color, and a, b, ca, and cb, which represent the remaining number of balls and distinct color count that can be distributed in each box.

The double nested loop helps in going through every possible way of allocating the balls of the current color to the boxes. For each ball, if the quantity A of that color in box A is less than or equal to the remaining capacity of box A, and if the quantity B of that color in box B is less than or equal to the remaining capacity of box B, it calculates the number of ways this can be done using the combination array (comb). It also increments the color count (ca and cb) if any balls of a new color were added to the box. The recursive call to dfs is multiplied by the number of ways this distribution can be done, and this is added to the result.

This recurses with reduced capacities and possibly incremented color counts until all balls are exhausted. If the capacities of the boxes are equal at that point and the color counts are also equal, we have a successful distribution and 1 is returned, else 0 is returned.

Finally, the result from the call to dfs is divided by the total number of ways the balls can be distributed between the boxes, regardless of the count of distinct colors, to get the probability.

In conclusion, by implementing recursive logic with some backing mathematics in combination with effective looping and conditions to ensure even distribution of balls, our Javascript solution above solves the problem.

Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

TA 👨‍🏫