Pythonで銀行の警備員からの最短距離を求める方法（BFS活用）

問題の概要

「O」「G」「W」の3種類の文字で構成されたマトリックス（二次元グリッド）を考えてみましょう。「O」は通路（開放スペース）、「G」は警備員、「W」は壁を表しています。この課題では、すべての「O」を、最も近い警備員からの最短距離に置き換えます。移動の際に壁は通過できません。最終的な出力では、警備員のセルは0、壁のセルは-1として表示されます。

入力例

出力例

解法のアプローチ：マルチソースBFS

この問題は、すべての警備員を起点として同時に探索を開始する「マルチソースBFS（幅優先探索）」を使うことで効率的に解くことができます。BFSは重みのないグラフにおいて最短距離を保証する探索手法であるため、各セルから最寄りの警備員までの距離を正確に求められます。

アルゴリズムの手順

1. マトリックスのサイズを表す変数MとNを5に設定します。
2. 上下左右への移動方向を表す配列 dir_row := [-1, 0, 1, 0]、dir_col := [0, 1, 0, -1] を用意します。
3. 関数 is_ok(i, j) を定義します。座標(i, j)がマトリックスの範囲外の場合はFalseを返します。
4. 関数 isSafe(i, j, matrix, result) を定義します。対象セルが「O」であり、かつ result[i][j] が未訪問（-1）の場合のみTrueを返します。
5. 関数 calculate_dist(matrix) を定義し、以下の処理を実行します。

calculate_dist() の内部処理

• M×Nの結果格納用マトリックス result を作成し、すべての要素を-1で初期化します。
• 両端キュー（deque）q を用意します。
• すべてのセルを走査し、警備員「G」が見つかったら、その位置と距離0をキューに追加し、result[i][j] を0に設定します。
• キューが空になるまで以下を繰り返します。
  • キューから要素を取り出し、座標(x, y)と現在の距離distを取得します。
  • 4方向それぞれについて、is_ok() と isSafe() の条件を満たす隣接セルがあれば、その距離を dist+1 に更新し、新しい位置情報をキューに追加します。
• 最後に、result マトリックスの全要素を出力します。

Pythonでの実装例

それでは、実際の実装を見てみましょう。

from collections import deque as queue

M = 5
N = 5
dir_row = [-1, 0, 1, 0]
dir_col = [0, 1, 0, -1]

# 座標(i, j)がマトリックスの範囲内かどうかを判定
def is_ok(i, j):
    if ((i < 0 or i > M - 1) or (j < 0 or j > N - 1)):
        return False
    return True

# セルが訪問可能（「O」かつ未設定）かどうかを判定
def isSafe(i, j, matrix, result):
    if (matrix[i][j] != 'O' or result[i][j] != -1):
        return False
    return True

def calculate_dist(matrix):
    # 結果用マトリックスを-1で初期化
    result = [[-1 for i in range(N)] for i in range(M)]
    q = queue()
    # すべての警備員を起点としてキューに追加
    for i in range(M):
        for j in range(N):
            if (matrix[i][j] == 'G'):
                pos = [i, j, 0]
                q.appendleft(pos)
                result[i][j] = 0
    # マルチソースBFSを実行
    while (len(q) > 0):
        curr = q.pop()
        x, y, dist = curr[0], curr[1], curr[2]
        for i in range(4):
            if is_ok(x + dir_row[i], y + dir_col[i]) and isSafe(x + dir_row[i], y + dir_col[i], matrix, result):
                result[x + dir_row[i]][y + dir_col[i]] = dist + 1
                pos = [x + dir_row[i], y + dir_col[i], dist + 1]
                q.appendleft(pos)
    # 結果を表示
    for i in range(M):
        for j in range(N):
            print(result[i][j], end=" ")
        print()

matrix = [['O', 'O', 'O', 'O', 'G'],
          ['O', 'O', 'O', 'W', 'O'],
          ['O', 'W', 'O', 'O', 'O'],
          ['G', 'W', 'W', 'W', 'O'],
          ['O', 'O', 'O', 'O', 'G']]

calculate_dist(matrix)

入力

[['O', 'O', 'O', 'O', 'G'],
['O', 'O', 'O', 'W', 'O'],
['O', 'W', 'O', 'O', 'O'],
['G', 'W', 'W', 'W', 'O'],
['O', 'O', 'O', 'O', 'G']]

出力

3 3 2 1 0
2 3 3 -1 1
1 -1 4 3 2
0 -1 -1 -1 1
1 2 2 1 0

計算量について

このアルゴリズムの時間計算量は O(M×N) です。各セルは最大でも1回しかキューに追加されないためです。空間計算量も O(M×N) となり、結果マトリックスとキューの保存に必要となります。壁で区切られた領域には到達できないため、その部分は自動的に-1のまま残る点にも注目してください。