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

• Time: $O((n / k)^k)$
• Space: $O(k)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 enum class BoxCase { kEqualBalls, kEqualDistantBalls }; class Solution { public: double getProbability(vector& balls) { const int n = accumulate(begin(balls), end(balls), 0) / 2; return cases(balls, 0, 0, 0, 0, 0, n, BoxCase::kEqualDistantBalls) / cases(balls, 0, 0, 0, 0, 0, n, BoxCase::kEqualBalls); } private: const vector fact{1, 1, 2, 6, 24, 120, 720}; // Assume we have two boxes A and B double cases(const vector& balls, int i, int ballsCountA, int ballsCountB, int colorsCountA, int colorsCountB, int n, BoxCase boxCase) { if (ballsCountA > n || ballsCountB > n) return 0; if (i == balls.size()) return boxCase == BoxCase::kEqualBalls ? 1 : colorsCountA == colorsCountB; double ans = 0; // Balls taken from A for balls[i] for (int ballsTakenA = 0; ballsTakenA <= balls[i]; ++ballsTakenA) { const int ballsTakenB = balls[i] - ballsTakenA; const int newcolorsCountA = colorsCountA + (ballsTakenA > 0); const int newcolorsCountB = colorsCountB + (ballsTakenB > 0); ans += cases(balls, i + 1, ballsCountA + ballsTakenA, ballsCountB + ballsTakenB, newcolorsCountA, newcolorsCountB, n, boxCase) / (fact[ballsTakenA] * fact[ballsTakenB]); } return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 enum BoxCase { kEqualBalls, kEqualDistantBalls } class Solution { public double getProbability(int[] balls) { final int n = Arrays.stream(balls).sum() / 2; return cases(balls, 0, 0, 0, 0, 0, n, BoxCase.kEqualDistantBalls) / cases(balls, 0, 0, 0, 0, 0, n, BoxCase.kEqualBalls); } private int[] fact = {1, 1, 2, 6, 24, 120, 720}; // Assume we have two boxes A and B double cases(int[] balls, int i, int ballsCountA, int ballsCountB, int colorsCountA, int colorsCountB, int n, BoxCase boxCase) { if (ballsCountA > n || ballsCountB > n) return 0; if (i == balls.length) return boxCase == BoxCase.kEqualBalls ? 1 : (colorsCountA == colorsCountB ? 1 : 0); double ans = 0; // Balls taken from A for balls[i] for (int ballsTakenA = 0; ballsTakenA <= balls[i]; ++ballsTakenA) { final int ballsTakenB = balls[i] - ballsTakenA; final int newcolorsCountA = colorsCountA + (ballsTakenA > 0 ? 1 : 0); final int newcolorsCountB = colorsCountB + (ballsTakenB > 0 ? 1 : 0); ans += cases(balls, i + 1, ballsCountA + ballsTakenA, ballsCountB + ballsTakenB, newcolorsCountA, newcolorsCountB, n, boxCase) / (fact[ballsTakenA] * fact[ballsTakenB]); } return ans; } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 from enum import Enum class BoxCase(Enum): kEqualDistantBalls = 0 kEqualBalls = 1 class Solution: def getProbability(self, balls: List[int]) -> float: n = sum(balls) // 2 fact = [1, 1, 2, 6, 24, 120, 720] def cases(i: int, ballsCountA: int, ballsCountB: int, colorsCountA: int, colorsCountB, boxCase: BoxCase) -> float: if ballsCountA > n or ballsCountB > n: return 0 if i == len(balls): return 1 if boxCase == BoxCase.kEqualBalls else colorsCountA == colorsCountB ans = 0.0 # Balls taken from A for balls[i] for ballsTakenA in range(balls[i] + 1): ballsTakenB = balls[i] - ballsTakenA newcolorsCountA = colorsCountA + (ballsTakenA > 0) newcolorsCountB = colorsCountB + (ballsTakenB > 0) ans += cases(i + 1, ballsCountA + ballsTakenA, ballsCountB + ballsTakenB, newcolorsCountA, newcolorsCountB, boxCase) / \ (fact[ballsTakenA] * fact[ballsTakenB]) return ans return cases(0, 0, 0, 0, 0, BoxCase.kEqualDistantBalls) / \ cases(0, 0, 0, 0, 0, BoxCase.kEqualBalls)