Pythonで2つの島を結ぶ最短の橋の長さを求めるプログラム（DFS＋BFS解説）

問題の概要

0を水、1を陸地とするバイナリ行列が与えられます。「島」とは、上下左右の4方向で連結している「1」の集合のことであり、各島は周囲を水（0）または行列の端に囲まれています。この問題では、2つの島を結ぶ最短の橋の長さを求めます。

例えば、次のような入力が与えられた場合を考えてみましょう。

この場合の出力は 1 となります。点 (1,0) と点 (1,2) を結ぶ橋を1つ架ければ、2つの島がつながるためです。

解法のアプローチ

この問題は、次の2段階の手法で効率的に解くことができます。

1. DFS（深さ優先探索）を使って、最初に見つかった島全体のセルを記録する。
2. BFS（幅優先探索）を使って、その島の周囲から水路を同心円状に広げ、別の島に到達した時点での距離を答えとする。

BFSを用いることで、必ず最短距離で反対側の島に到達できることが保証されます。

具体的な手順

• 行列の行数を row、列数を col として取得します。
• 関数 dfs(i, j, s) を定義します。
  • (i, j) がすでに集合 s に含まれている場合は何もせず戻ります。
  • mat[i][j] が 0（水）の場合も戻ります。
  • (i, j) を s に追加します。
  • 範囲内である限り、上下左右の4方向に対して再帰的に dfs を呼び出します。
• メイン処理では、まず空の集合 seen を用意します。
• 行列を走査し、最初に見つかった「1」のセルから dfs を実行して、1つ目の島全体を seen に登録します。
• 両端キュー（deque）q を用意します。
• seen 内の各陸地セルについて、隣接するセルが水（0）であれば、その座標と距離 1 を q に追加します。
• q が空になるまで以下を繰り返します。
  • (i, j, dist) をキューの左端から取り出します。
  • (i, j) がすでに seen に含まれていればスキップします。
  • (i, j) を seen に追加します。
  • mat[i][j] が 1（陸地）であれば、dist − 1 を返します。これは、最後の1歩で島に乗ったため、橋の長さからその分を差し引く必要があるからです。
  • 範囲内の隣接セルすべてについて、距離 dist + 1 として q に追加します。

実装例

それでは、実際のPythonコードを見てみましょう。

import collections
class Solution:
    def solve(self, mat):
        row = len(mat)
        col = len(mat[0])
        def dfs(i, j, s):
            if (i, j) in s:
                return
            if mat[i][j] == 0:
                return
            s.add((i, j))
            if i - 1 >= 0:
                dfs(i - 1, j, s)
            if i + 1 < row:
                dfs(i + 1, j, s)
            if j - 1 >= 0:
                dfs(i, j - 1, s)
            if j + 1 < col:
                dfs(i, j + 1, s)
        seen = set()
        for i in range(row):
            if len(seen) > 0:
                break
            for j in range(col):
                if mat[i][j] == 1:
                    dfs(i, j, seen)
                    break
        q = collections.deque()
        for land in seen:
            i, j = land
            if i - 1 >= 0 and mat[i - 1][j] == 0:
                q.append((i - 1, j, 1))
            if i + 1 < row and mat[i + 1][j] == 0:
                q.append((i + 1, j, 1))
            if j - 1 >= 0 and mat[i][j - 1] == 0:
                q.append((i, j - 1, 1))
            if j + 1 < col and mat[i][j + 1] == 0:
                q.append((i, j + 1, 1))
        while len(q) > 0:
            i, j, dist = q.popleft()
            if (i, j) in seen:
                continue
            seen.add((i, j))
            if mat[i][j] == 1:
                return dist - 1
            if i - 1 >= 0:
                q.append((i - 1, j, dist + 1))
            if i + 1 < row:
                q.append((i + 1, j, dist + 1))
            if j - 1 >= 0:
                q.append((i, j - 1, dist + 1))
            if j + 1 < col:
                q.append((i, j + 1, dist + 1))
ob = Solution()
matrix = [
    [0, 0, 1],
    [1, 0, 1],
    [1, 0, 0],
]
print(ob.solve(matrix))

入力

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

出力

1

計算量について

このアルゴリズムでは、DFSとBFSのそれぞれで各セルを高々1回ずつ訪問するため、時間計算量は O(row × col)、使用するメモリも訪問管理のための集合やキューによって O(row × col) となります。グリッド系の探索問題としては非常に効率的な解法です。