Pythonでk回の移動後にインデックス0へ戻る経路の数を求めるプログラム

問題概要

長さ n のリストの位置 0(インデックス 0)からスタートします。各ステップでは、「右に 1 つ進む」「左に 1 つ戻る(リストの範囲外には出られない)」「その場にとどまる」のいずれかの操作を選べます。ここで、ちょうど k ステップを使って再びインデックス 0 に戻ってくる一意な移動経路(ウォーク)の総数を求めるのがこの問題です。答えが非常に大きくなる可能性があるため、10^9 + 7 で割った余りを返します。

たとえば、入力が n = 7、k = 4 の場合、出力は 9 になります。条件を満たす操作列は次の 9 通りです。

• [右、右、左、左]
• [右、左、右、左]
• [待機、待機、待機、待機]
• [右、左、待機、待機]
• [待機、待機、右、左]
• [右、待機、待機、左]
• [右、待機、左、待機]
• [待機、右、左、待機]
• [待機、右、待機、左]

解き方のアプローチ

この問題は、再帰と動的計画法(DP)の考え方を組み合わせることで解けます。手順は以下のとおりです。

• m := 10^9 + 7(剰余を取るための定数)
• N := リストの長さ
• dp(i, jumps) という関数を定義する。i は現在位置、jumps は残りのステップ数
• jumps が 0 のとき、i が 0 なら 1、そうでなければ 0 を返す
• count := dp(i, jumps − 1)(その場にとどまるケース)
• i ≥ 0 の場合、count := count + dp(i + 1, jumps − 1)(右へ移動するケース)
• i ≤ N − 1 の場合、count := count + dp(i − 1, jumps − 1)(左へ移動するケース)
• count を返す
• メイン処理では dp(0, k) mod m を返す

それでは、実際の実装を見て理解を深めましょう。

実装例(Python)

class Solution:
   def solve(self, length, n):
      MOD = 10 ** 9 + 7
      N = length

      def dp(i, jumps):
         if jumps == 0:
            return +(i == 0)

         count = dp(i, jumps - 1)
         if i >= 0:
            count += dp(i + 1, jumps - 1)
         if i <= N - 1:
            count += dp(i - 1, jumps - 1)
         return count
      return dp(0, n) % MOD    
ob = Solution()
n = 7
k = 4
print(ob.solve(n, k))

入力

7, 4

出力

9

補足:計算量について

上記の実装は素朴な再帰のみで構成されているため、ステップ数 k が大きくなると計算量が指数的に増大します。functools.lru_cache などを使ってメモ化を行えば、状態数は「現在位置 × 残りステップ数」に抑えられるため、大幅な高速化が可能です。競技プログラミングなどで大きな入力を扱う場合は、メモ化の導入を検討するとよいでしょう。