Python演算法：如何解決樓梯臺階問題

本文由碼農網 – 小峰原創翻譯，轉載請看清文末的轉載要求，歡迎參與我們的付費投稿計劃！

讓我們考慮以下問題。

有一個有N個臺階的樓梯，你一次可以爬1或2個臺階。

給定N，編寫一個函式，返回爬完樓梯的方式數量。步驟的順序很重要。

例如，如果N是4，那麼有5種方式：

• 1,1,1,1
• 2,1,1
• 1,2,1
• 1,1,2
• 2,2

如果規定的不是一次只能爬1或2步，而是可以使用正整數X集合內的任意數字爬樓梯，那會怎麼樣？例如，如果X = {1,3,5}，則表示一次爬升1，3或5階樓梯。

解決方案

從一些測試案例開始總是好的做法。讓我們從小的案例開始，看看能否找到某種規律。

• N = 1，1種爬樓方式：[1]
• N = 2，2種爬樓方式：[1,1]，[2]
• N = 3，3種爬樓方式：[1,2]，[1,1,1]，[2,1]
• N = 4，5種爬樓方式：[1,1,2]，[2,2]，[1,2,1]，[1,1,1,1]，[2,1,1]

你有沒有注意到什麼？請看N = 3時，爬完3階樓梯的方法數量是3，基於N = 1和N = 2。存在什麼關係？

爬完N = 3的兩種方法是首先達到N = 1，然後再往上爬2步，或達到N = 2再向上爬1步。所以 f(3) = f(2) + f(1)。

這對N = 4是否成立呢？是的，這也是成立的。因為我們只能在達到第三個臺階然後再爬一步，或者在到了第二個臺階之後再爬兩步這兩種方式爬完4個臺階。所以f(4) = f(3) + f(2)。

所以關係如下： f(n) = f(n – 1) + f(n – 2)，且f(1) = 1和f(2) = 2。這就是斐波那契數列。

def fibonacci(n):
    if n <= 1:
        return 1
    return fibonacci(n - 1) + fibonacci(n - 2)

當然，這很慢（O(2^N)）——我們要做很多重複的計算！通過迭代計算，我們可以更快：

def fibonacci(n):
    a, b = 1, 2
    for _ in range(n - 1):
        a, b = b, a + b
    return a

現在，讓我們嘗試概括我們學到的東西，看看是否可以應用到從集合X中取步數這個要求下的爬樓梯。類似的推理告訴我們，如果X = {1,3,5}，那麼我們的演算法應該是f(n) = f(n – 1) + f(n – 3) + f(n – 5)。如果n <0，那麼我們應該返回0，因為我們不能爬負數。

def staircase(n, X):
    if n < 0:
        return 0
    elif n == 0:
        return 1
    elif n in X:
        return 1 + sum(staircase(n - x, X) for x in X if x < n)
    else:
        return sum(staircase(n - x, X) for x in X if x < n)

這也很慢（O(|X|^N)），因為也重複計算了。我們可以使用動態程式設計來加快速度。

每次的輸入cache[i]將包含我們可以用集合X到達臺階i的方法的數量。然後，我們將使用與之前相同的遞迴從零開始構建陣列：

def staircase(n, X):
    cache = [0 for _ in range(n + 1)]
    cache[0] = 1
    for i in range(n + 1):
        cache[i] += sum(cache[i - x] for x in X if i - x > 0)
        cache[i] += 1 if i in X else 0
    return cache[-1]

現在時間複雜度為O(N * |X|)，空間複雜度為O(N)。

歡迎繼續探索其他有趣的程式設計問題。