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

// Accepted solution for LeetCode #1467: Probability of a Two Boxes Having The Same Number of Distinct Balls
class Solution {
    private int n;
    private long[][] c;
    private int[] balls;
    private Map<List<Integer>, Long> f = new HashMap<>();

    public double getProbability(int[] balls) {
        int mx = 0;
        for (int x : balls) {
            n += x;
            mx = Math.max(mx, x);
        }
        n >>= 1;
        this.balls = balls;
        int m = Math.max(mx, n << 1);
        c = new long[m + 1][m + 1];
        for (int i = 0; i <= m; ++i) {
            c[i][0] = 1;
            for (int j = 1; j <= i; ++j) {
                c[i][j] = c[i - 1][j - 1] + c[i - 1][j];
            }
        }
        return dfs(0, n, 0) * 1.0 / c[n << 1][n];
    }

    private long dfs(int i, int j, int diff) {
        if (i >= balls.length) {
            return j == 0 && diff == 0 ? 1 : 0;
        }
        if (j < 0) {
            return 0;
        }
        List<Integer> key = List.of(i, j, diff);
        if (f.containsKey(key)) {
            return f.get(key);
        }
        long ans = 0;
        for (int x = 0; x <= balls[i]; ++x) {
            int y = x == balls[i] ? 1 : (x == 0 ? -1 : 0);
            ans += dfs(i + 1, j - x, diff + y) * c[balls[i]][x];
        }
        f.put(key, ans);
        return ans;
    }
}

// Accepted solution for LeetCode #1467: Probability of a Two Boxes Having The Same Number of Distinct Balls
func getProbability(balls []int) float64 {
	n, mx := 0, 0
	for _, x := range balls {
		n += x
		mx = max(mx, x)
	}
	n >>= 1
	m := max(mx, n<<1)
	c := make([][]int, m+1)
	for i := range c {
		c[i] = make([]int, m+1)
	}
	for i := 0; i <= m; i++ {
		c[i][0] = 1
		for j := 1; j <= i; j++ {
			c[i][j] = c[i-1][j-1] + c[i-1][j]
		}
	}
	k := len(balls)
	f := make([][][]int, k)
	for i := range f {
		f[i] = make([][]int, n+1)
		for j := range f[i] {
			f[i][j] = make([]int, k<<1|1)
			for h := range f[i][j] {
				f[i][j][h] = -1
			}
		}
	}
	var dfs func(int, int, int) int
	dfs = func(i, j, diff int) int {
		if i >= k {
			if j == 0 && diff == k {
				return 1
			}
			return 0
		}
		if j < 0 {
			return 0
		}
		if f[i][j][diff] != -1 {
			return f[i][j][diff]
		}
		ans := 0
		for x := 0; x <= balls[i]; x++ {
			y := 1
			if x != balls[i] {
				if x == 0 {
					y = -1
				} else {
					y = 0
				}
			}
			ans += dfs(i+1, j-x, diff+y) * c[balls[i]][x]
		}
		f[i][j][diff] = ans
		return ans
	}
	return float64(dfs(0, n, k)) / float64(c[n<<1][n])
}

# Accepted solution for LeetCode #1467: Probability of a Two Boxes Having The Same Number of Distinct Balls
class Solution:
    def getProbability(self, balls: List[int]) -> float:
        @cache
        def dfs(i: int, j: int, diff: int) -> float:
            if i >= k:
                return 1 if j == 0 and diff == 0 else 0
            if j < 0:
                return 0
            ans = 0
            for x in range(balls[i] + 1):
                y = 1 if x == balls[i] else (-1 if x == 0 else 0)
                ans += dfs(i + 1, j - x, diff + y) * comb(balls[i], x)
            return ans

        n = sum(balls) >> 1
        k = len(balls)
        return dfs(0, n, 0) / comb(n << 1, n)

// Accepted solution for LeetCode #1467: Probability of a Two Boxes Having The Same Number of Distinct Balls
/**
 * [1467] 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).
 * 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.
 *  
 * 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 equal 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
 *
 *  
 * Constraints:
 *
 * 	1 <= balls.length <= 8
 * 	1 <= balls[i] <= 6
 * 	sum(balls) is even.
 *
 */
pub struct Solution {}

// problem: https://leetcode.com/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls/
// discuss: https://leetcode.com/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls/discuss/?currentPage=1&orderBy=most_votes&query=

// submission codes start here

impl Solution {
    // Credit: https://leetcode.com/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls/solutions/3159815/just-a-runnable-solution/
    pub fn get_probability(balls: Vec<i32>) -> f64 {
        let balls = balls.iter().map(|&x| x as i64).collect::<Vec<_>>();
        let sum = balls.iter().sum::<_>();
        let mut a = vec![0; balls.len()];
        let mut b = vec![0; balls.len()];
        Solution::dfs_helper(&sum, &balls, &mut a, &mut b, 0, 0, 0) / Solution::perm(&balls)
    }

    fn perm(a: &[i64]) -> f64 {
        let mut result = 1.0;
        let mut j = 1;
        for i in 0..a.len() {
            for k in 1..=a[i] {
                result = result * j as f64 / k as f64;
                j += 1;
            }
        }
        result
    }

    fn dfs_helper(
        sum: &i64,
        balls: &[i64],
        a: &mut [i64],
        b: &mut [i64],
        i: usize,
        sa: i64,
        sb: i64,
    ) -> f64 {
        if sa > sum / 2 || sb > sum / 2 {
            return 0.0;
        }
        if i == balls.len() {
            let ca = a.iter().filter(|&&x| x > 0).count();
            let cb = b.iter().filter(|&&x| x > 0).count();
            if ca != cb {
                return 0.0;
            }
            return Solution::perm(a) * Solution::perm(b);
        }
        let mut result = 0.0;
        for j in 0..=balls[i] {
            a[i] = j;
            b[i] = balls[i] - j;
            result += Solution::dfs_helper(sum, balls, a, b, i + 1, sa + a[i], sb + b[i]);
        }
        result
    }
}

// submission codes end

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_1467_example_1() {
        let balls = vec![1, 1];

        let result = 1.00000;

        assert_f64_near!(Solution::get_probability(balls), result);
    }

    #[test]
    fn test_1467_example_2() {
        let balls = vec![2, 1, 1];

        let result = 0.6666666666666667;

        assert_f64_near!(Solution::get_probability(balls), result);
    }

    #[test]
    fn test_1467_example_3() {
        let balls = vec![1, 2, 1, 2];

        let result = 0.60000;

        assert_f64_near!(Solution::get_probability(balls), result);
    }
}

// Accepted solution for LeetCode #1467: Probability of a Two Boxes Having The Same Number of Distinct Balls
function getProbability(balls: number[]): number {
    const n = balls.reduce((a, b) => a + b, 0) >> 1;
    const mx = Math.max(...balls);
    const m = Math.max(mx, n << 1);
    const c: number[][] = Array(m + 1)
        .fill(0)
        .map(() => Array(m + 1).fill(0));
    for (let i = 0; i <= m; ++i) {
        c[i][0] = 1;
        for (let j = 1; j <= i; ++j) {
            c[i][j] = c[i - 1][j - 1] + c[i - 1][j];
        }
    }
    const k = balls.length;
    const f: number[][][] = Array(k)
        .fill(0)
        .map(() =>
            Array(n + 1)
                .fill(0)
                .map(() => Array((k << 1) | 1).fill(-1)),
        );
    const dfs = (i: number, j: number, diff: number): number => {
        if (i >= k) {
            return j === 0 && diff === k ? 1 : 0;
        }
        if (j < 0) {
            return 0;
        }
        if (f[i][j][diff] !== -1) {
            return f[i][j][diff];
        }
        let ans = 0;
        for (let x = 0; x <= balls[i]; ++x) {
            const y = x === balls[i] ? 1 : x === 0 ? -1 : 0;
            ans += dfs(i + 1, j - x, diff + y) * c[balls[i]][x];
        }
        return (f[i][j][diff] = ans);
    };
    return dfs(0, n, k) / c[n << 1][n];
}