Pythonで幅優先探索を実装する

Python

アルゴリズムの勉強のために、幅優先探索を書いてみた。
使ったのはAtCoder Beginers Contest 007Cの問題。この頃はアルゴリズムがそのまま出題されてたようだ。

特殊事項として、この問題ではスタートからゴールまでは必ず行くことができる前提がある。さらに周り中が壁で囲まれているので、盤面からはみ出すのを考慮する必要がない。

実装

from collections import deque

def bfs(maze, visited, sy, sx, gy, gx):
    queue = deque([[sy, sx]])
    visited[sy][sx] = 0
    while queue:
        y, x = queue.popleft()
        if [y, x] == [gy, gx]:
            return visited[y][x]
        for j, k in ([1, 0], [-1, 0], [0, 1], [0, -1]):
            new_y, new_x = y+j, x+k
            if maze[new_y][new_x] == "." and \
                visited[new_y][new_x] == -1:
                visited[new_y][new_x] = visited[y][x] + 1
                queue.append([new_y, new_x])

if __name__ == "__main__":
    R, C = map(int, input().split())
    sy, sx = map(int, input().split())
    gy, gx = map(int, input().split())
    sy, sx, gy, gx = sy-1, sx-1, gy-1, gx-1

    maze = [input() for i in range(R)]
    visited = [[-1]*C for j in range(R)]
    print(bfs(maze, visited, sy, sx, gy, gx))

たった25行で書けた。

原理としては、今いるマスから上下左右に隣り合うマスの座標をキューに入れるとともに、visitedリストにも記録していくという繰り返し。
キューは先入先出なので、今いるマスから上下左右に隣り合うマスを全て探索した後、それらの上下左右に隣り合うマスを探索する、というように探索を行っていく。
全マス調べるため時間はかかるが、ゴールへの最短ルートを確実に求められる。

速度アップの工夫

実をいうと、上記のコードは2度目の提出だった。
その前に書いたバージョンでは、bfs関数を以下のようにしていた。

def bfs(maze, visited, sy, sx, gy, gx):
    queue = deque([[sy, sx]])
    count = 0
    while queue:
        #ここでforループを回している
        for _ in range(len(queue)):
            coord = queue.popleft()
            visited[coord[0]][coord[1]] = 1
            if coord == [gy, gx]:
                return count
            for j, k in ([1, 0], [-1, 0], [0, 1], [0, -1]):
                next_coord = [coord[0]+j, coord[1]+k]
                if maze[next_coord[0]][next_coord[1]] == "." and \
                   visited[next_coord[0]][next_coord[1]] != 1:
                    queue.append(next_coord)
        count += 1

visitedリストの値を0=未訪問と1=訪問のみとしていたため、bfs関数の中のwhile queue以下でforループを回して、ループを何回まわったかをカウントしスタートからの距離を求めていた。しかし、テストケースの最初の4分の1くらいで実行時間が2秒を超え、あえなくTLEに。
そこで、visitedリストの値にスタートからの距離を直接書かせることにした。まずスタートマスの値を0、未訪問のマスの値を-1としておく。そして新しいマスを探索するときは、visitedリストから前にいたマスの値を取り出し、それに1を足した値を新たに書けばよい。
こうすることで、forループをひとつ減らすことができ、実行時間も全体で25msと大幅に速くなった。リストをうまく使うことで、回数カウント目的のforループをなくすというテクニックは、他の問題でも応用できそうだ。