Problem Description

You are given a set of k different colors of balls. The i-th color has balls[i] balls. All balls of the same color are indistinguishable, but balls of different colors are distinguishable.

Your task is to divide all the balls into two boxes such that:

• Each box gets at least one ball.
• The probability that both boxes have the same number of distinct colors is calculated.

Constraints:

• 1 <= balls.length <= 8
• 1 <= balls[i] <= 6
• sum(balls) <= 48

Thought Process

At first glance, this problem appears to be a classic combinatorial partitioning problem. Since balls of the same color are indistinguishable, it's about distributing counts rather than permutations. The brute-force way is to try every possible way to split the balls and count the number of ways that result in both boxes having the same number of distinct colors.

However, as the number of colors and balls increases, the number of ways grows exponentially. Optimizing means avoiding recomputation by leveraging recursion, memoization, or combinatorial mathematics. We need to count the number of valid distributions (where the number of distinct colors in each box is equal) and divide it by the total number of ways to distribute the balls.

The key insight is to use recursion to try all possible splits for each color, keeping track of how many colors are present in each box, and use combinatorics to count arrangements efficiently.

Solution Approach

Here’s how we solve the problem step by step:

1. Recursive Distribution: For each color, decide how many balls of that color go into box 1 (the rest go to box 2). For example, if there are 3 red balls, you can put 0, 1, 2, or 3 in box 1.
2. Tracking State: As we recurse through colors, keep track of:
   • The number of balls in each box
   • The number of distinct colors in each box (a color is present if at least one ball of that color is in the box)
3. Base Case: When all colors have been processed, check if both boxes have the same number of distinct colors and at least one ball.
4. Counting Ways: For each distribution, the number of ways to choose which balls go to box 1 is the product of binomial coefficients for each color: C(balls[i], x) where x is the number of balls of color i in box 1.
5. Probability Calculation: The answer is the total number of valid distributions (where both boxes have the same number of distinct colors) divided by the total number of ways to distribute the balls.
6. Optimization: Use memoization and precompute factorials for fast binomial coefficient calculation.

This approach ensures we only count unique distributions (not permutations), and efficiently explores all possibilities using recursion.

Example Walkthrough

Example Input: balls = [2, 1] (2 balls of color A, 1 ball of color B)

Total balls: 3. We want to split them into two boxes, each with at least one ball, and count the ways both boxes have the same number of distinct colors.

• Possible splits for color A (2 balls): 0, 1, or 2 in box 1 (rest in box 2)
• Possible splits for color B (1 ball): 0 or 1 in box 1 (rest in box 2)

Let's enumerate all possible distributions:

• (A:0, B:0) → Invalid (box 1 empty)
• (A:0, B:1) → Box 1: 1B, Box 2: 2A. Distinct colors: Box 1=1, Box 2=1. Valid
• (A:1, B:0) → Box 1: 1A, Box 2: 1A 1B. Distinct colors: Box 1=1, Box 2=2. Invalid
• (A:1, B:1) → Box 1: 1A 1B, Box 2: 1A. Distinct colors: Box 1=2, Box 2=1. Invalid
• (A:2, B:0) → Box 1: 2A, Box 2: 1B. Distinct colors: Box 1=1, Box 2=1. Valid
• (A:2, B:1) → Invalid (box 2 empty)

Ways to choose balls:

• For (A:0, B:1): C(2,0)*C(1,1) = 1*1 = 1
• For (A:2, B:0): C(2,2)*C(1,0) = 1*1 = 1

Total possible arrangements:

• For each color, split between 0 and total balls for box 1 (excluding all in one box):
• All combinations: (2+1)*(1+1) = 3*2 = 6
• But exclude cases where one box is empty: (A:0,B:0) and (A:2,B:1)
• So, total = 4

Time and Space Complexity

Brute-force approach:

• For each color i, try all possible splits (0 to balls[i])
• Total number of recursive calls: product(balls[i]+1)
• For k colors and maximum n balls per color, this is O((n+1)^k)
• Space: O(k) for recursion stack

• Same as above, but with memoization, redundant states are avoided
• Precomputing factorials for fast combinations: O(N) space for factorials, where N is total balls
• Overall, still exponential in k (number of colors), but feasible due to small constraints (k ≤ 8)

Summary

This problem is a classic example of combinatorial distribution with constraints. The elegant solution uses recursion to explore all possible ball splits, tracks the number of distinct colors in each box, and efficiently counts arrangements using binomial coefficients. By dividing the number of valid distributions by the total number of distributions, we compute the required probability. The approach leverages the small input constraints and combinatorial mathematics for an efficient and clear solution.

Code Implementation

from math import comb
from functools import lru_cache

class Solution:
    def getProbability(self, balls):
        n = len(balls)
        total = sum(balls)
        half = total // 2
        
        # Precompute factorials up to 48 for combinations
        max_balls = sum(balls)
        fact = [1]
        for i in range(1, max_balls+1):
            fact.append(fact[-1]*i)
        def C(n, k):
            return fact[n] // (fact[k]*fact[n-k])

        total_ways = 0
        valid_ways = 0

        def dfs(i, box1, box2, color1, color2, ways):
            nonlocal total_ways, valid_ways
            if i == n:
                if box1 == 0 or box2 == 0: return
                if box1 != box2: return
                if color1 == color2:
                    valid_ways += ways
                total_ways += ways
                return
            for pick in range(balls[i]+1):
                next_box1 = box1 + pick
                next_box2 = box2 + (balls[i] - pick)
                next_color1 = color1 + (1 if pick > 0 else 0)
                next_color2 = color2 + (1 if balls[i] - pick > 0 else 0)
                dfs(i+1, next_box1, next_box2, next_color1, next_color2, ways * C(balls[i], pick))
        dfs(0, 0, 0, 0, 0, 1)
        return valid_ways / total_ways
      

#include <vector>
#include <numeric>
#include <algorithm>
using namespace std;

class Solution {
public:
    vector fact;
    long long C(int n, int k) {
        return fact[n] / (fact[k] * fact[n-k]);
    }
    double getProbability(vector<int>& balls) {
        int n = balls.size();
        int total = accumulate(balls.begin(), balls.end(), 0);
        int max_balls = total;
        fact = vector<long long>(max_balls+1, 1);
        for (int i = 1; i <= max_balls; ++i)
            fact[i] = fact[i-1] * i;
        long long valid_ways = 0, total_ways = 0;
        function<void(int,int,int,int,int,long long)> dfs = [&](int i, int box1, int box2, int color1, int color2, long long ways) {
            if (i == n) {
                if (box1 == 0 || box2 == 0) return;
                if (box1 != box2) return;
                if (color1 == color2)
                    valid_ways += ways;
                total_ways += ways;
                return;
            }
            for (int pick = 0; pick <= balls[i]; ++pick) {
                int next_box1 = box1 + pick;
                int next_box2 = box2 + (balls[i] - pick);
                int next_color1 = color1 + (pick > 0 ? 1 : 0);
                int next_color2 = color2 + ((balls[i] - pick) > 0 ? 1 : 0);
                dfs(i+1, next_box1, next_box2, next_color1, next_color2, ways * C(balls[i], pick));
            }
        };
        dfs(0, 0, 0, 0, 0, 1LL);
        return (double)valid_ways / total_ways;
    }
};
      

import java.util.*;

class Solution {
    long[] fact;
    long C(int n, int k) {
        return fact[n] / (fact[k] * fact[n-k]);
    }
    double validWays = 0, totalWays = 0;
    public double getProbability(int[] balls) {
        int n = balls.length;
        int total = 0;
        for (int b : balls) total += b;
        fact = new long[total+1];
        fact[0] = 1;
        for (int i = 1; i <= total; ++i)
            fact[i] = fact[i-1] * i;
        dfs(balls, 0, 0, 0, 0, 0, 1L);
        return validWays / totalWays;
    }
    void dfs(int[] balls, int i, int box1, int box2, int color1, int color2, long ways) {
        if (i == balls.length) {
            if (box1 == 0 || box2 == 0) return;
            if (box1 != box2) return;
            if (color1 == color2)
                validWays += ways;
            totalWays += ways;
            return;
        }
        for (int pick = 0; pick <= balls[i]; ++pick) {
            int nextBox1 = box1 + pick;
            int nextBox2 = box2 + (balls[i] - pick);
            int nextColor1 = color1 + (pick > 0 ? 1 : 0);
            int nextColor2 = color2 + ((balls[i] - pick) > 0 ? 1 : 0);
            dfs(balls, i+1, nextBox1, nextBox2, nextColor1, nextColor2, ways * C(balls[i], pick));
        }
    }
}
      

var getProbability = function(balls) {
    const n = balls.length;
    const total = balls.reduce((a,b) => a+b, 0);
    // Precompute factorials
    let fact = [1];
    for (let i = 1; i <= total; ++i)
        fact[i] = fact[i-1] * i;
    function C(n, k) {
        return fact[n] / (fact[k]*fact[n-k]);
    }
    let validWays = 0, totalWays = 0;
    function dfs(i, box1, box2, color1, color2, ways) {
        if (i === n) {
            if (box1 === 0 || box2 === 0) return;
            if (box1 !== box2) return;
            if (color1 === color2)
                validWays += ways;
            totalWays += ways;
            return;
        }
        for (let pick = 0; pick <= balls[i]; ++pick) {
            let nextBox1 = box1 + pick;
            let nextBox2 = box2 + (balls[i] - pick);
            let nextColor1 = color1 + (pick > 0 ? 1 : 0);
            let nextColor2 = color2 + ((balls[i] - pick) > 0 ? 1 : 0);
            dfs(i+1, nextBox1, nextBox2, nextColor1, nextColor2, ways * C(balls[i], pick));
        }
    }
    dfs(0, 0, 0, 0, 0, 1);
    return validWays / totalWays;
};