りんごとバナナとエンジニア

エンジニア修行の記録

Pythonで深さ優先探索を実装する

前の記事で幅優先探索を実装したので、ついでにこちらも。
問題はAtCoder Typical Contest 001A。一応解説スライドはあるが、スタックを使った実装方法が書いてないので、自分でやってみることにした。

実装

from collections import deque

def dfs(maze, visited, sh, sw):
    stack = deque([[sh, sw]])
    visited[sh][sw] = 1
    while stack:
        h, w = stack.pop()
        if maze[h][w] == "g":
            return "Yes"
        for j, k in ([1, 0], [-1, 0], [0, 1], [0, -1]):
            new_h, new_w = h+j, w+k
            if new_h < 0 or new_h >= H or \
               new_w < 0 or new_w >= W:
                continue
            elif maze[new_h][new_w] != "#" and \
                visited[new_h][new_w] == 0:
                visited[new_h][new_w] = 1
                stack.append([new_h, new_w])
    
    return "No"

if __name__ == "__main__":
    H, W = map(int, input().split())
    maze = [input() for _ in range(H)]

    for i in range(H):
        if maze[i].find("s") != -1:
            sh = i
            sw = maze[i].find("s")
            break

    visited = [[0]*W for _ in range(H)]
    print(dfs(maze, visited, sh, sw))

基本的には、幅優先探索の場合と流れはほぼ同じ。
ただし、今回は後入先出のスタックを使っているため、上下左右にマスが分岐する際に、どれかひとつをこれ以上進めなくなるまで調べ尽くす。

今回の問題のように、単にゴールまでのルートがあるかを素早く調べるには、深さ優先探索の方がよいだろう。