2020年6月3日 星期三

Leetcode題解 Python:Probability of a Two Boxes Having The Same Number of Distinct Balls

有 2n 個球,球有 k 個顏色,寫成一串 k 長度 balls 陣列,balls[i] 代表第 i 顏色有多少顆球。

今天要隨機分到兩箱子,分完打開看,回傳最後兩箱子球的顏色數目相同的機率。

因為是隨機,所以拿到顏色順序不會按照顏色種類排列,例:[1,3] [3,1] ,而二者都是合規則的可能情形。 

按照真實隨機分配的方式,即一顆顆隨機抽出,運算次數會隨 k 與 sum(balls) 呈指數爆炸性成長。
如:[6,6,6,6,6,6],會作 36! 次探討,自然會超時,要想辦法改進。

對於隨機排列組合,像是抽籤等情況,要掌握到這是「邊抽邊開」還是「抽完再開」?二者在計算上是不同的。
「抽完再開」的如此題,同顏色的順序就不重要,如:白1白3、白3白1,都視作 白白,視為同一種。

如果有6顆球白白白白紅紅,那抽完再看的組合數有 6!/4!/2! 個,因同顏色間排列都視為一個,所以用除去 4!(白) 2!(紅) 來排除同色順序的影響。

為了要加速計算,勢必得抛棄直觀的一顆顆隨機分配,朝著直接抽一把起來的方向分配會快個多。
照順序一次次抽,這裡的順序選擇「顏色」,從第 0 位抽到第 k-1 位顏色,最後再把組合數 n! / C0balls! / C1balls! / ... / Ck-1balls! 算出來。

拿 [6,6,6,6,6,6] 為例,在 [0] 位顏色時, box1 可拿 0-6 顆,算作 7 種拿法,那麼最糟要探討次數為 7**6 ,感覺很多,但相比 36!,
算快上非常顯著的。

接下來有二種DP:
1.紀錄BOX1的各色球數,最後可反推BOX2。 (i, tuple([B1C0balls, B1C1balls, ..., B1Ck-1balls]))
2.紀錄BOX1與BOX2的目前的球數與顏色數。 (i, color1, balls1, factorial1, color2, balls2, factorial2)

我用 2. 。 (2. 是 1. 的攤平化,用tuple是非常沒有效率的)

寫遞迴函數時,先思考中止條件,也就是分完了 k 色。此時算組合數的 n! 是已知,但後面的除數是上層分球時決定的,
因此最深層的DFS需要把除數一路往深傳遞。

算出 box1 組合數,再乘上 box2 組合數,就能得到該分配下的隨機總可能數,納入總數。要是 c1 == c2,就納入合格數。

回傳 合格數  /  總數

Python
class Solution:
    def getProbability(self, balls: List[int]) -> float:
        n = sum(balls) // 2
        bs = len(balls)
        
        @lru_cache(None)
        def fl(num):
            if num <= 1: return 1
            return num * fl(num - 1)

        @lru_cache(None)
        def dfs(i, c1, n1, f1, c2, n2, f2):
            if i == bs:                
                combo = fl(n)*2 / f1 / f2
                return combo * (c1 == c2), combo
            maxTake = min(balls[i], n - n1)
            minTake = max(balls[i] - (n - n2), 0)
            s, p = 0, 0
            for b in range(minTake, maxTake + 1):
                res = dfs(i+1, c1+(1 if b > 0 else 0), n1+b, f1*fl(b), c2+(1 if b < balls[i] else 0), n2 + balls[i] - b, f2*fl(balls[i] - b))
                s += res[0]
                p += res[1]
            return s, p
     
        return dfs(0, 0, 0, 1, 0, 0, 1)[0]  / dfs(0, 0, 0, 1, 0, 0, 1)[1]