Pythonで目的地に到達するために必要な高さの増加量の最小値を求めるアルゴリズム

問題の概要

各セルの高さを格納した行列 M が与えられます。ここで M[r][c] はセル (r, c) の高さを表します。現在、左上隅に位置しており、右下隅へ移動したいと考えています。隣接するセル（上下左右）へは、そのセルの高さが現在いるセルの高さ「以下」である場合にのみ移動できます。ただし、移動を開始する前に、好きなだけ多くのセルの高さを上げることが可能です。このとき、右下のセルに到達するために必要な「高さの増加量の合計」の最小値を求めるのがこの問題です。

入力例

この場合、答えは 4 になります。経路 [2, 4, 5, 1] をたどることを考え、途中のセルの高さを次のように変更すればよいからです。

この変更では、値 2 → 5、4 → 5、5 → 5 と上げており、増加量の合計は 3 + 1 + 0 = 4 となります。

解き方のアプローチ

この問題は、ダイクストラ法（Dijkstra's algorithm）を応用することで効率的に解けます。ポイントは、右下のゴールから左上のスタートへ向かって逆方向に探索することです。経路上のセルの高さは単調に減少していく必要があるため、ゴール側から「そのセルで維持すべき高さ」と「高さを上げるためにかかったコストの累積」を状態として管理しながら探索を進めます。

具体的な手順は以下の通りです。

• INF（無限大）を定義しておきます。

• R, C にそれぞれ行列の行数・列数を設定します。

• ヒープによる優先度キュー pq を作成し、初期状態として [0, R-1, C-1, M[-1][-1]]（コスト 0、右下の座標、右下の元の高さ）を挿入します。

• 最短コストを記録する辞書 dist を用意し、dist[(R-1, C-1, 右下の高さ)] = 0 とします。

• pq が空になるまで以下を繰り返します。

  • pq から最小の要素を取り出し、d（コスト）、r（行）、c（列）、h（高さ）に格納します。

  • dist[(r, c, h)] より d の方が大きい場合は、既により良い状態が記録済みなので次の反復へ進みます。

  • r と c がどちらも 0（左上）に到達していれば、d を答えとして返します。

  • 隣接セル [r+1, c]、[r, c+1]、[r-1, c]、[r, c-1] のそれぞれについて、行列の範囲内であれば次を処理します。

    • h2 = max(隣接セルの元の高さ, h) として、隣接セルに必要な高さを求めます。

    • d2 = d + max(h2 − 隣接セルの元の高さ, 0) として、新しいコストを計算します。

    • d2 が dist[(nr, nc, h2)] より小さければ、dist を更新し、[d2, nr, nc, h2] を pq に挿入します。

実装例（Python）

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

import collections
import heapq

class Solution:
    def solve(self, A):
        INF = float('inf')
        R, C = len(A), len(A[0])

        # 右下のゴールから逆方向にダイクストラ法を適用
        pq = [[0, R-1, C-1, A[-1][-1]]]
        dist = collections.defaultdict(lambda: INF)
        dist[R-1, C-1, A[-1][-1]] = 0

        while pq:
            d, r, c, h = heapq.heappop(pq)
            # 既により良いコストが記録されている場合はスキップ
            if dist[r, c, h] < d:
                continue
            # 左上に到達したらコストを返す
            if r == c == 0:
                return d
            for nr, nc in [[r+1, c], [r, c+1], [r-1, c], [r, c-1]]:
                if 0 <= nr < R and 0 <= nc < C:
                    h2 = max(A[nr][nc], h)
                    d2 = d + max(h2 - A[nr][nc], 0)
                    if d2 < dist[nr, nc, h2]:
                        dist[nr, nc, h2] = d2
                        heapq.heappush(pq, [d2, nr, nc, h2])

ob = Solution()
matrix = [
    [2, 4, 5],
    [8, 6, 1]
]
print(ob.solve(matrix))

入力

[[2, 4, 5],[8, 6, 1]]

出力

4

まとめ

このように、ゴールからスタートへ逆向きにダイクストラ法を適用することで、「通過するセルの高さが単調非増加になる」という移動条件を満たしながら、最小限の高さ増加で目的地に到達する経路を見つけることができます。優先度キューを活用することで、座標と高さの組み合わせからなる状態空間を効率的に探索でき、実用的な計算量で最適解を得られます。