【Python】3段登りが最大k回まで使える階段問題を動的計画法で解く方法

n段の階段と整数kが与えられた場面を考えてみましょう。最初は0段目に立っており、一度の移動で1段、2段、または3段ずつ上ることができます。ただし、「3段ずつの移動」は全体で最大k回までしか使えません。この条件下で、階段を一番上まで登る方法が何通りあるかを求めるのがこの問題です。

たとえば n = 5、k = 2 の場合、答えは 13 になります。実際にあり得る登り方は次の13通りです。

• [1, 1, 1, 1, 1]
• [2, 1, 1, 1]
• [1, 2, 1, 1]
• [1, 1, 2, 1]
• [1, 1, 1, 2]
• [1, 2, 2]
• [2, 1, 2]
• [2, 2, 1]
• [1, 1, 3]
• [1, 3, 1]
• [3, 1, 1]
• [2, 3]
• [3, 2]

解き方のアプローチ（動的計画法）

この問題は動的計画法（DP）を使うことで効率よく解けます。鍵となるのは、次のような二次元のメモ化テーブルを定義することです。

memo[j][i] ＝ 「3段登りを最大j回まで使えるとき、i段を登る方法の総数」

このテーブルを以下の手順で埋めていきます。

1. n = 0 の場合は 1 を返します（何もしないという1通り）。
2. n = 1 の場合も 1 を返します。
3. k = min(k, n) としておきます（n段しかなければ、それ以上の3段登りは意味がないため）。
4. (n+1) × (k+1) のサイズの二次元リスト memo を用意します。
5. すべての行 r について、初期値として memo[r][0] = 1、memo[r][1] = 1、memo[r][2] = 2 を設定します。
6. 3段登りを一切使わない行（j = 0）は、フィボナッチ数列と同じ漸化式になります。
   memo[0][i] = memo[0][i-1] + memo[0][i-2]
7. j ≥ 1 の各行について、i = 3 から n まで次のように計算します。まず count = i // 3（i段を登るときに考えられる3段登りの最大回数）を求めます。
   • count ≤ j の場合（3段登りの回数制約が実質的に効かない場合）：
     memo[j][i] = memo[j][i-1] + memo[j][i-2] + memo[j][i-3]
   • count > j の場合（ここで3段登りを使うと残り回数が1減る）：
     memo[j][i] = memo[j][i-1] + memo[j][i-2] + memo[j-1][i-3]
8. 最後に memo[k][n] を返せば答えとなります。

ポイントは、直前の一手が1段・2段・3段のどれであるかで場合分けしている点です。「3段登り」で到達する場合は使用回数の予算を1消費するため、参照先が memo[j-1][i-3] に変わります。

Pythonでの実装例

以下が実際のコードです。

class Solution:
    def solve(self, n, k):
        if n == 0:
            return 1
        if n == 1:
            return 1
        k = min(k, n)
        memo = [[0] * (n + 1) for _ in range(k + 1)]
        for r in range(k + 1):
            memo[r][0] = 1
            memo[r][1] = 1
            memo[r][2] = 2
        for i in range(3, n + 1):
            memo[0][i] = memo[0][i - 1] + memo[0][i - 2]
        for j in range(1, k + 1):
            for i in range(3, n + 1):
                count = i // 3
                if count <= j:
                    memo[j][i] = memo[j][i - 1] + memo[j][i - 2] + memo[j][i - 3]
                else:
                    memo[j][i] = memo[j][i - 1] + memo[j][i - 2] + memo[j - 1][i - 3]
        return memo[k][n]


ob = Solution()
print(ob.solve(n=5, k=2))

入力

n = 5, k = 2

出力

13

計算量について

このアルゴリズムでは、テーブルの各セルを高々定数回の加算で埋めていくため、時間計算量は O(n × k)、必要なメモリも同じく O(n × k) となります。全探索で毎回組み合わせを数え上げる方法に比べ、大幅に高速に動作するのが動的計画法の利点です。