pep.py

from functools import lru_cache

class Solution:
    def getProbability(self, A):
        def binom(n, k):
            if not(0 <= k <= n):
                return 0
            ans = 1
            for i in range(1, k + 1):
                ans *= n - i + 1
                ans //= i
            return ans
        
        N = sum(A) // 2

        @lru_cache(None)
        def dp(i, left, right, bal):
            if left < 0 or right < 0:
                return 0.0
            if i == len(A):
                return 1.0 if bal == 0 else 0.0

            ans = 0.0            
            for k1 in range(A[i] + 1):
                k2 = A[i] - k1
                if k1 <= left and k2 <= right:
                    prob = binom(left, k1) * binom(right, k2)
                    prob /= binom(left + right, k1 + k2)
                    b = -1 if k1 == 0 else 1 if k2 == 0 else 0
                    ans += prob * dp(i + 1, left - k1, right - k2, bal + b)
            
            return ans
        
        return dp(0, N, N, 0)