Pythonで開始点から終了点までコストkのパスの数をカウントするプログラム

問題の概要

0と1のみで構成される2次元のバイナリ行列と、ある値 k が与えられます。左上のセルからスタートし、右下のセルへ向かいます。1回の移動で進めるのは「下」または「右」のいずれか一方のみです。パスのスコアは、そのパスが通過するセルの値の合計として定義されます。このとき、開始セルから終了セルまでのパスの中で、スコアがちょうど k に等しくなるものの総数を求めます。パスの候補数が非常に大きくなる可能性があるため、その場合は結果を 10⁹+7 で割った余りを返すこととします。

入力例

k = 2 の場合、出力は 4 になります。スコアが 2 となるパスは [R,R,D,D]、[D,R,R,D]、[D,D,R,R]、[D,R,D,R] の4通りです(※ D は下方向、R は右方向への移動を表します)。

解法のアプローチ

この問題は深さ優先探索(DFS)を用いて解くことができます。各セルに到達した時点での累積スコアを管理しながら、「下」と「右」の2方向へ再帰的に探索を進めていきます。具体的な手順は以下の通りです。

• deno := 10^9 + 7(剰余を取るための定数)

• m := 行列の行数、n := 行列の列数

• dfs(i, j, pts) 関数を定義する(i は現在の行、j は現在の列、pts は累積スコア)

• i >= m または j >= n の場合(範囲外に出た場合)、0 を返す

• pts := pts + matrix[i][j](現在のセルの値を加算)

• i = m - 1 かつ j = n - 1 の場合(右下のセルに到達した場合)、pts が k と等しければ 1、そうでなければ 0 を返す

• それ以外の場合は、dfs(i + 1, j, pts) + dfs(i, j + 1, pts) を返す(下と右の両方向を探索)

• メイン処理では、dfs(0, 0, 0) mod deno を返す

実装例(Python)

以下のコードで実際の動作を確認できます。

class Solution:
    def solve(self, matrix, k):
        m, n = len(matrix), len(matrix[0])
        def dfs(i=0, j=0, pts=0):
            if i >= m or j >= n:
                return 0
            pts += matrix[i][j]
            if i == m - 1 and j == n - 1:
                return int(pts == k)
            return dfs(i + 1, j, pts) + dfs(i, j + 1, pts)
        return dfs() % (10 ** 9 + 7)

ob = Solution()
matrix = [
    [0, 0, 1],
    [1, 0, 1],
    [0, 1, 0]
]
k = 2
print(ob.solve(matrix, k))

入力

[
    [0, 0, 1],
    [1, 0, 1],
    [0, 1, 0]
], 2

出力

4

補足:計算量について

上記の素朴な DFS は、存在するパスの本数に比例して計算時間が増加するため、行列が大きくなると非現実的な速度になります。実務では、メモ化(動的計画法)を組み合わせて状態 (i, j, pts) ごとの結果をキャッシュすることで、大幅な高速化が可能です。また、1つのパスが通過するセル数は m + n - 1 個であり、各セルの値は 0 か 1 なので、累積スコア pts の最大値は m + n - 1 以下になります。この性質を利用すると、探索空間をさらに絞り込むこともできます。