Pythonで2次元行列から収集できるコインの最大量を求めるプログラム

問題の概要

ここでは、2次元の行列(マトリックス)を扱います。各セル mat[r][c] には、そのマスに置かれたコインの枚数が格納されています。プレイヤーは任意のマスからスタートし、「上・下・左・右」の4方向(斜め移動は不可)に移動しながらコインを集めていきます。

移動したマスのコインはすべて回収され、そのマスの値は 0 に変わります。さらに、コインが 0 枚のマスには立ち入ることができないという制約があります。この条件下で、収集できるコインの合計の最大値を求めるのが本記事の目的です。

これは典型的なグラフ探索の問題であり、一度訪れたマスは二度と使えないため、「重み付きグリッド上の最長経路問題」として捉えることができます。解法としては、深さ優先探索(DFS) と バックトラッキング を組み合わせるのが有効です。

入力例と出力

例として、次のような行列を考えてみましょう。

この場合、出力は 18 になります。2 → 3 → 6 → 4 → 3 という経路をたどることで、コインを合計18枚収集できるためです。なお、右下の 12 は周囲が 0 のマスに囲まれているため、そもそも到達できません。

アルゴリズムの手順

この問題を解くために、以下の手順に従います。

• 行列が空の場合は 0 を返します。
• 行数を n、列数を m とします。
• 移動方向を表すオフセット x = [-1, 1, 0, 0]、y = [0, 0, -1, 1] を用意します(それぞれ上下左右に対応)。
• 再帰関数 util(a, b) を定義します。この関数は、マス (a, b) から探索を始めたときに収集できるコインの最大値を返します。
  • ret = 0 で初期化します。
  • 4方向それぞれについて、移動先 (t1, t2) が有効なマス(範囲内かつコインが残っている)であれば:
    • そのマスのコイン数 t を退避させ、mat[t1][t2] = 0 として訪問済みマークを付けます。
    • ret = max(ret, util(t1, t2) + t) で最大値を更新します。
    • 探索が終わったら mat[t1][t2] = t と元に戻します(バックトラッキング)。
  • 最終的な ret を返します。
• メイン処理では、すべてのマスを走査し、コインが存在するマスを起点として util() を呼び出します。起点のマス自体も一時的に 0 にしてから呼び出し、得られた結果の最大値を res に記録して返します。

Pythonでの実装例

class Solution:
   def solve(self, mat):
      if not mat:
         return 0
      n, m = len(mat), len(mat[0])
      x, y = [-1, 1, 0, 0], [0, 0, -1, 1]

      def ok(a, b):
         # 範囲内かつコインが残っているマスかどうかを判定
         return 0 <= a < n and 0 <= b < m and mat[a][b]

      def util(a, b):
         ret = 0
         for k in range(4):
            t1, t2 = x[k] + a, y[k] + b
            if ok(t1, t2):
               t, mat[t1][t2] = mat[t1][t2], 0  # 訪問済みにする
               ret = max(ret, util(t1, t2) + t)
               mat[t1][t2] = t  # バックトラック(状態を元に戻す)
         return ret

      res = 0
      for i in range(n):
         for j in range(m):
            if mat[i][j]:
               temp, mat[i][j] = mat[i][j], 0  # 起点も訪問済みにする
               res = max(res, util(i, j) + temp)
      return res

ob = Solution()
matrix = [
   [2, 4, 3],
   [3, 6, 0],
   [2, 0, 12]
]
print(ob.solve(matrix))

入力

[
   [2, 4, 3],
   [3, 6, 0],
   [2, 0, 12]
]

出力

18

計算量に関する注意点

各マスから最大4方向へ分岐する再帰探索を行うため、最悪の場合、計算量はマスの総数に対して指数的に増大します(O(4^n×m) オーダー)。そのため、このアプローチは小〜中規模のグリッドに適しています。より大きな入力を扱う場合は、訪問済みマスの集合をビット列で管理するビットDP(動的計画法)などによる高速化を検討するとよいでしょう。