1
class Solution:
2
    def getProbability(self, balls: List[int]) -> float:
3
        m = len(balls)
4
        N = sum(balls)
5
        n = N // 2
6

7
        prefix = [0] * m
8
        for i in range(m):
9
            prefix[i] = prefix[i - 1] + balls[i]
10

11
        # STEP 1: Compute the number of ways to pick j balls from i total, i.e. C(i, j)
12
        choose = [[0] * (N + 1) for i in range(N + 1)]
13
        choose[0][0] = 1
14
        for i in range(1, N + 1):
15
            for pick in range(N + 1):
16
                # DECISION 1: don't pick the ith ball
17
                choose[i][pick] += choose[i - 1][pick]
18
                # DECISION 2: pick the ith ball
19
                choose[i][pick] += choose[i - 1][pick - 1]
20

21
        # STEP 2: From first i ball types, compute ways to:
22
        # - pick c1 balls in box1
23
        # - such that the difference in unique ball count between box1 and box2 is d
24
        ways = [[defaultdict(int) for k in range(n + 1)] for i in range(m + 1)]
25
        ways[0][0][0] = 1
26

27
        for i in range(m):
28
            b = balls[i]
29
            if i == 0:
30
                prev_total = 0
31
            else:
32
                prev_total = prefix[i - 1]
33
            for c1 in range(n + 1):
34
                c2 = prev_total - c1
35
                if c2 < 0:
36
                    continue
37
                for d in ways[i][c1]:
38
                    for add1 in range(b + 1):
39
                        add2 = b - add1
40
                        if c1 + add1 > n:
41
                            continue
42
                        if c2 + add2 > n:
43
                            continue
44
                        if add1 == b:
45
                            delta = 1
46
                        elif add2 == b:
47
                            delta = -1
48
                        else:
49
                            delta = 0
50
                        ways_to_add = choose[b][add1] * ways[i][c1][d]
51
                        ways[i + 1][c1 + add1][d + delta] += ways_to_add
52

53
        # compute the actual probability
54
        return ways[m][n][0] / choose[N][n]