Probability of a two boxes having the same number of distinct balls

Array Math Dynamic Programming Backtracking Combinatorics Probability and Statistics

61.0% acceptance

Feb 25, 2026

297

178

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. Return the probability that the two boxes have the same number of distinct balls. Answers within 10^-5 of the actual value will be accepted as correct.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn get_probability(balls: Vec<i32>) -> f64 {
    let _k = balls.len();
    let total: i32 = balls.iter().sum();
    let n = total / 2;

    // Precompute factorials for combinations
    let mut fact = vec![1.0f64; 49];
    for i in 1..49 { fact[i] = fact[i - 1] * i as f64; }
    let comb = |a: i32, b: i32| -> f64 {
      if b < 0 || b > a { 0.0 } else { fact[a as usize] / (fact[b as usize] * fact[(a - b) as usize]) }
    };

    let mut favorable = 0.0f64;
    let mut denominator = 0.0f64;

    // Backtracking: for each color, pick how many go to box1
    fn bt(
      color: usize, balls: &[i32], n: i32,
      taken: i32, dist1: i32, dist2: i32,
      weight: f64, favorable: &mut f64, denominator: &mut f64,
      comb: &dyn Fn(i32, i32) -> f64,
    ) {
      if taken > n { return; }
      if color == balls.len() {
        if taken != n { return; }
        *denominator += weight;
        if dist1 == dist2 { *favorable += weight; }
        return;
      }
      let _remaining_colors = (balls.len() - color) as i32;
      let remaining_balls: i32 = balls[color..].iter().sum();
      // pruning: can we still reach taken==n?
      if taken + remaining_balls < n { return; }
      if taken > n { return; }

      for x in 0..=balls[color] {
        let new_dist1 = dist1 + if x > 0 { 1 } else { 0 };
        let new_dist2 = dist2 + if x < balls[color] { 1 } else { 0 };
        let new_weight = weight * comb(balls[color], x);
        bt(color + 1, balls, n, taken + x, new_dist1, new_dist2, new_weight, favorable, denominator, comb);
      }
    }

    bt(0, &balls, n, 0, 0, 0, 1.0, &mut favorable, &mut denominator, &comb);
    favorable / denominator
  }
}