https://leetcode.com/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls

This is an interesting mathematical problem. Denote by k different number of balls and by B1, ..., Bk number of balls of each type. Then we need to put into first box X1, ..., Xk balls and into the second B1 - X1, ..., Bk - Xk. We need to check 2 properties:

  1. Number of balls in two boxes are equal.
  2. Number of ball types in each boxes are equal.

Then we need to evaluate total number of possibilites for boxes, using Multinomial coefficients, (see https://en.wikipedia.org/wiki/Multinomial_theorem for more details) and divide it by total number of possibilities. Let us work through example: B1, B2, B3 = 1, 2, 3, then all possible ways to put into boxes are:

  1. (0, 0, 0) and (1, 2, 3)
  2. (0, 0, 1) and (1, 2, 2)
  3. (0, 0, 2) and (1, 2, 1)

  1. (1, 2, 1) and (0, 0, 2)
  2. (1, 2, 2) and (0, 0, 1)
  3. (1, 2, 3) and (0, 0, 0)

But only for 2 of them holds the above two properties:

  1. (0, 2, 1) and (1, 0, 2)
  2. (1, 0, 2) and (0, 2, 1)

So, the answer in this case will be [M(0, 2, 1) * M(1, 0, 2) + M(1, 0, 2) * M(0, 2, 1)]/ M(1, 2, 3)

Complexity We check all possible (B1+1) * ... * (Bk+1) splits into two boxes, and then evaluate number of zeros and multinomial coefficientes, so it is O(B1 ... Bk * k)

Update There is more elegant way to use multinomial coefficients, which I updated.

class Solution:
    def multinomial(self, n):
        return factorial(sum(n))/prod([factorial(i) for i in n])
  
    def getProbability(self, balls):
        k, n, Q = len(balls), sum(balls)// 2, 0
        arrays = [range(0,i+1) for i in balls]
        t = list(product(*arrays))
        for i in range(len(t)):
            if sum(t[i]) == n and t[i].count(0) == t[-i-1].count(0):
                Q += self.multinomial(t[i]) * self.multinomial(t[-i-1]) 

        return Q / self.multinomial(list(balls))     

If you like the solution, you can upvote it on leetcode discussion section: Problem 1467