Pythonで2つの場所にある金塊を回収する最小コストを求めるプログラム

問題の概要

2次元の行列(グリッド)と、開始位置・目標位置を表す複数の値(row、col、erow0、ecol0、erow1、ecol1)が与えられます。現在位置は matrix[row][col] であり、matrix[erow0][ecol0] と matrix[erow1][ecol1] の2箇所に置かれた金塊を回収したいとします。

移動は上下左右の4方向が可能ですが、セル (r, c) に立ち入るときにはコスト matrix[r][c] を支払う必要があります。ただし、同じセルに何度足を踏み入れても、そのセルのコストは最初の1回しか支払いません。この条件下で、2つの金塊をどちらも回収するときの最小コストを求めるのがこの問題の目的です。

入力例と出力の確認

たとえば、次のような入力を考えてみます。

row = 0、col = 0、erow0 = 0、ecol0 = 3、erow1 = 2、ecol1 = 2 の場合、出力は 8 になります。これは、(0, 0) を起点として (0, 3) と (2, 2) の金塊を回収するケースです。まず (0, 0) から (0, 3) まで3ステップで移動し、その後 (0, 0) に戻って、値が 1 のセルをたどるルートで (2, 2) へ向かうのが最適解となります。

解法のアプローチ:ダイクストラ法

この問題は、各セルをノード、上下左右の隣接セルへの移動をエッジとみなしたグラフの最短経路問題として捉えることができます。そこで、ダイクストラ法を3回実行する方針を採ります。

1. スタート地点 (row, col) から全セルへの最小コスト a を求める
2. 金塊1の位置 (erow0, ecol0) から全セルへの最小コスト b を求める
3. 金塊2の位置 (erow1, ecol1) から全セルへの最小コスト c を求める
4. すべてのセル (i, j) を「スタートと2つの金塊をつなぐ合流地点」と仮定し、a[i][j] + b[i][j] + c[i][j] − 2 × matrix[i][j] の最小値を答えとする

合流地点のコストを2回分差し引くのは、3つのダイクストラ計算のそれぞれで合流地点セルのコストが重複してカウントされるためです。「同じセルの再訪では追加コストが発生しない」というルール上、合流地点で支払うべきコストは実質1回分だけでよいことになります。

アルゴリズムの詳細な手順

• 関数 is_valid(x, y) を定義します。x と y が行列の範囲内なら true、範囲外なら false を返します。
• 関数 min_cost(sx, sy) を定義します。起点 (sx, sy) からダイクストラ法を実行し、全セルへの最小コスト表を返します。
  • heap に (matrix[sx][sy], sx, sy) を初期要素として格納します。
  • dists は元の行列と同じサイズで、すべて無限大(math.inf)で初期化した行列です。
  • dists[sx][sy] に matrix[sx][sy] を設定します。
  • heap が空になるまで、次の処理を繰り返します。
    • (cost, x, y) を heap の先頭から取り出します。
    • 隣接セル候補 [(x, y−1)、(x+1, y)、(x−1, y)、(x, y+1)] のそれぞれ (nx, ny) について、is_valid(nx, ny) が真 かつ matrix[nx][ny] + cost < dists[nx][ny] である場合:
      • edge := matrix[nx][ny]
      • dists[nx][ny] := edge + cost
      • (edge + cost, nx, ny) を heap に挿入します。
  • 最後に dists を返します。
• メイン処理では、res を無限大で初期化し、a := min_cost(row, col)、b := min_cost(erow0, ecol0)、c := min_cost(erow1, ecol1) を計算します。
• すべてのセル (i, j) について、res と a[i][j] + b[i][j] + c[i][j] − 2 × matrix[i][j] の小さいほうを res に代入します。
• 最終的な res を返します。

Pythonでの実装例

理解を深めるために、以下の実装を見てみましょう。

import heapq
import math

class Solution:
    def solve(self, matrix, row, col, erow0, ecol0, erow1, ecol1):
        def is_valid(x, y):
            return x >= 0 and y >= 0 and x < len(matrix) and y < len(matrix[0])

        def min_cost(sx, sy):
            heap = [(matrix[sx][sy], sx, sy)]
            dists = [[math.inf] * len(matrix[0]) for _ in range(len(matrix))]
            dists[sx][sy] = matrix[sx][sy]
            while heap:
                cost, x, y = heapq.heappop(heap)
                for nx, ny in [(x, y - 1), (x + 1, y), (x - 1, y), (x, y + 1)]:
                    if is_valid(nx, ny) and matrix[nx][ny] + cost < dists[nx][ny]:
                        edge = matrix[nx][ny]
                        dists[nx][ny] = edge + cost
                        heapq.heappush(heap, (edge + cost, nx, ny))
            return dists

        res = math.inf
        a, b, c = min_cost(row, col), min_cost(erow0, ecol0), min_cost(erow1, ecol1)
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                res = min(res, a[i][j] + b[i][j] + c[i][j] - 2 * matrix[i][j])
        return res

ob = Solution()
matrix = [
    [1, 1, 1, 1, 1],
    [1, 10, 10, 10, 10],
    [1, 1, 1, 10, 10]
]
row = 0
col = 0
erow0 = 0
ecol0 = 3
erow1 = 2
ecol1 = 2
print(ob.solve(matrix, row, col, erow0, ecol0, erow1, ecol1))

入力

[
  [1, 1, 1, 1, 1],
  [1, 10, 10, 10, 10],
  [1, 1, 1, 10, 10]
], 0, 0, 0, 3, 2, 2

出力

8

計算量について

グリッドのサイズを N 行 M 列とすると、ダイクストラ法を1回実行するのに O(NM log(NM)) の時間がかかります。今回は3回実行するため、全体の時間計算量も O(NM log(NM)) となります(定数倍は無視)。また、距離表を3つ保持するため、空間計算量は O(NM) です。