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

• Time: $O((n / k)^k)$
• Space: $O(k)$

C++JavaPython

enum class BoxCase { kEqualBalls, kEqualDistantBalls };

class Solution {
 public:
  double getProbability(vector<int>& balls) {
    const int n = accumulate(balls.begin(), balls.end(), 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<int> fact{1, 1, 2, 6, 24, 120, 720};

  // Assume we have two boxes A and B.
  double cases(const vector<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.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;
  }
};

enum BoxCase { EQUAL_BALLS, EQUAL_DISTANT_BALLS }

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.EQUAL_DISTANT_BALLS) /
        cases(balls, 0, 0, 0, 0, 0, n, BoxCase.EQUAL_BALLS);
  }

  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.EQUAL_BALLS ? 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;
  }
}

from enum import Enum


class BoxCase(Enum):
  EQUAL_DISTANT_BALLS = 0
  EQUAL_BALLS = 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.EQUAL_BALLS
                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.EQUAL_DISTANT_BALLS) /
            cases(0, 0, 0, 0, 0, BoxCase.EQUAL_BALLS))