雙蛋問題的 Python 遞歸解決

雙蛋問題的 Python 遞歸解決

今天看了 李永樂老師關於雙蛋問題的講解視頻，受用很大。本着好記性不如爛筆頭的精神，把這個問題記錄在此。

據傳某大廠有這樣一個面試題：手里有 2 個雞蛋，另外有 100 層樓。有一未知的臨界樓層，雞蛋從臨界樓層以下扔下去，一定不會碎；從臨界樓層以上丟下去，一定會碎。沒有摔碎的雞蛋可以反復使用，碎了的雞蛋就不能再往下扔了。問，在最糟糕的情況下，至少需要多少次能夠找到臨界樓層？

吐槽一句，這個雞蛋可能比較特殊，因為普通雞蛋別說 100 層樓，從桌子上掉下去基本就碎了。不過問題本身是很有價值的，我們可以把雞蛋改成玻璃球之類的，低樓層摔不碎，高樓層受不了就成了。

另外，要想讀懂本文，恐怕需要一點遞歸和算法基礎，否則不一定能看懂。這不是因為我水平高，寫的東西高深莫測。而是因為我的水平太低，道行太淺，目前還沒辦法做到深入淺出，十分抱歉。

好了，閑話不多扯，來談談解決問題的思路。

二分查找解決無限多雞蛋的情況

直接看問題，似乎沒什么思路。我們不妨稍微簡化一下問題。比如，如果我們有無數多個雞蛋，最壞的情況下至少需要幾個雞蛋能找到臨界樓層？

這就很簡單了，使用二分法即可。先從 50 層試，如果雞蛋碎了，說明臨界樓層在下面，就去 25 層再試；如果雞蛋沒碎，說明臨界樓層在上面，到 75 層去試。以此類推，每次排除一半的可能，很快就能找到答案。

二分法需要的次數的公式是 \(log_2(100)\) 向下取整再加 1，計算結果應該是 7。

遞歸法解決雙蛋問題

不過二分查找似乎並沒有對我們解決問題有什么特別好的啟發，我們只好另辟蹊徑。我們可不可以通過 分而治之 的思想來解決這個問題呢？

首先，基線條件很好確定：

1. 在有 2 個雞蛋的情況下，如果只有一層樓，只需要試一次；如果有兩層樓，只需要試兩次；如果沒有樓，那就干脆不用試了（看似是廢話，但是是很重要的邊界條件）。
2. 如果只有 1 個雞蛋，只能老老實實從下往上嘗試，也就是在最壞的情況下，有幾層樓就要試幾次。

接下來，我們就要思考遞歸條件了。如何能將問題簡化。

令在有 2 個雞蛋時，最壞的情況下，N 層樓所需要嘗試的最少次數為 \(T_N\)。

假設總共有 N 層樓，我們在第 K 層樓進行一次嘗試。那么此時，就會分成兩種情況：

1. 雞蛋在 K 層碎掉了，也就說明臨界樓層在 K 層以下。但是此時，我們只剩下 1 個雞蛋，最壞的情況下還要檢測 \(K - 1\) 次才能找到臨界樓層
2. 雞蛋在 K 層沒有碎，臨界樓層在 K 層以上。此時我們還是有 2 個雞蛋，還剩下 \(N-K\) 層樓需要檢測，那么最壞的情況下，還需要檢測 \(T_{N-K}\) 次。很顯然 \(N-K\) 要比 N 少，我們順利實現對問題的簡化。

最壞的情況顯然是 \(K - 1\) 和 \(T_{N-K}\) 兩個數的最大的那一個再加上 1，因為我們先試了一次。這個最大的數，就是 \(T_N\)。

不過這里面有一個 K 是不能確定的。為了找到合適的 K，我們需要把 K 從 1 到 N 的情況全部計算出來，找到使得 \(T_N\) 最小的情況即可。

用代碼來解決這個問題就是：

def two_egg(n: int) -> int:
    """
    雙蛋問題的遞歸求解
    :param n: 樓層數
    :return: 最壞情況下，找到臨界樓層所需最少嘗試次數
    """
    if n == 0:    # 沒有樓就不需要試
        return 0
    elif n == 1:   # 有一層樓，試一次
        return 1
    result_list = []
    for k in range(1, n + 1):    # 在每一層都試一下
        result_list.append(max(k - 1, two_egg(n - k)) + 1)    # 把每一層的情況都記錄下來
    return min(result_list)    # 最好的結果就是我們想要的


# 用 1 到 11 的數字測試，不用 100 是因為電腦性能不夠，測到 11 是因為 10 和 11 的結果不同
for f in range(1, 12):
    print(f'{f} -------> {two_egg(f)}')

上面的代碼用到了遞歸。隨着遞歸層數的增加，會占用很多資源，計算時間也會特別長。可以通過記錄低樓層的結果，優化上面的代碼：

def two_egg_opt(n: int, result_dict: dict) -> int:
    if n in result_dict:
        return result_dict[n]
    else:
        result_list = []
        for k in range(1, n + 1):  # 在每一層都試一下
            result_list.append(max(k - 1, two_egg_opt(n - k, result_dict)) + 1)  # 把每一層的情況都記錄下來
        result_dict[n] = min(result_list)  # 最好的結果就是我們想要的
        return min(result_list)

# 從前計算的結果記錄在result_dict中，下次使用可以直接拿，極大減少了遞歸層數
result_dict = {0: 0, 1: 1}
for i in range(1, 101):
    result_dict[i] = two_egg_opt(i, result_dict)
print(result_dict)

優化前的代碼用我的小電腦根本無法求出 100 層樓的雙蛋問題的解。而使用這個優化后的代碼，1 到 100 層樓雙蛋問題的解幾乎立刻就出來了。

遞歸法解決普遍雙蛋問題

用二分查找，可以解決雞蛋數目不限的情況，遞歸查找可以解決只有 2 個雞蛋的情況。現在，我們把問題進一步擴展：如果我們有 M 個雞蛋，N 層樓，在最壞的情況下，至少需要測試多少次能夠找到臨界樓層？

基線條件根上面的差不多一樣：

1. 不管有多少個雞蛋，如果只有一層樓，只需要試一次；如果沒有樓，那就干脆不用試了。
2. 如果只有 1 個雞蛋，只能老老實實從下往上嘗試，也就是在最壞的情況下，有幾層樓就要試幾次。

遞歸條件其實也很類似，只是因為雞蛋數目的引入，會稍微復雜一丁丁點點。

令在有 M 個雞蛋時，最壞的情況下，N 層樓所需要嘗試的最少次數為 \(T_{M,\space N}\)。

依舊假設總共有 N 層樓，我們在第 K 層樓進行一次嘗試。那么此時，還是會分成兩種情況：

1. 雞蛋在 K 層碎掉了，也就說明臨界樓層在 K 層以下。但是此時，我們只剩下 \(M-1\) 個雞蛋，最壞的情況下還要檢測 \(T_{M-1,\space K - 1}\) 次才能找到臨界樓層
2. 雞蛋在 K 層沒有碎，臨界樓層在 K 層以上。此時我們還是有 M 個雞蛋，還剩下 \(N-K\) 層樓需要檢測，那么最壞的情況下，還需要檢測 \(T_{M,\space N-K}\) 次

上面的兩種情況，要么簡化了雞蛋數量，要么簡化了樓層數量，最終都可以通過遞歸來找到答案。最終的結果需要是 \(T_{M-1,\space K - 1}\) 和 \(T_{M,\space N-K}\) 這兩個數中最大的那一個加上 1，因為我們最開始的時候在 K 層測試了一下。

同樣地，我們需要遍歷測試當 K 為 1 到 N 時的各種情況，取其中所需步驟最少的，就是我們要的結果。

用代碼表示就是：

def two_egg_general(m: int, n: int) -> int:
    """
    普遍雙蛋問題的解決
    :param m: 雞蛋數量
    :param n: 樓層總層數
    :return: 最糟糕的情況下，找到臨界樓層所需最少嘗試數目
    """
    if n == 0:    # 如果沒有樓，不需要試
        return 0
    elif n == 1:    # 只有 1 層樓，試一次就足夠
        return 1
    if m == 1:    # 只有 1 個蛋，有幾層樓就要使幾次
        return n
    result_list = []
    for k in range(1, n + 1):
        result_list.append(max(two_egg_general(m - 1, k - 1), two_egg_general(m, n - k)) + 1)
    return min(result_list)


for i in range(1, 12):
    for j in range(1, 12):
        print(f'({i}, {j}) --> {two_egg_general(i, j)}', end=' | ')
    print()

測試結果如下：

附上雙蛋問題的參照表，都是吻合的。只不過我是以樓層數為橫軸，雞蛋數為縱軸了而已。

同樣地，也可以對這個代碼進行優化：

def two_egg_gen_opt(m: int, n: int, result_dict: dict) -> int:
    """
    普遍雙蛋問題遞歸解決的優化
    :param m: 雞蛋數量
    :param n: 樓層總層數
    :param result_dict: 儲存結果的字典
    :return: 最糟糕的情況下，找到臨界樓層所需最少嘗試數目
    """
    if (m, n) in result_dict:
        return result_dict[(m, n)]
    if n == 0:    # 如果沒有樓，不需要試
        result_dict[(m, n)] = 0
        return 0
    elif n == 1:    # 只有 1 層樓，試一次就足夠
        result_dict[(m, n)] = 1
        return 1
    if m == 1:    # 只有 1 個蛋，有幾層樓就要使幾次
        result_dict[(m, n)] = n
        return n
    result_list = []
    for k in range(1, n + 1):
        result_list.append(max(two_egg_gen_opt(m - 1, k - 1, result_dict), two_egg_gen_opt(m, n - k, result_dict)) + 1)
    result_dict[(m, n)] = min(result_list)
    return min(result_list)


result_dict = {}
for i in range(1, 20):
    for j in range(1, 1002):
        print(f'({i}, {j}) --> {two_egg_gen_opt(i, j, result_dict)}', end=' | ')
    print()