dfsやbfsをpythonで実装してみる

(全文無料です)

競プロでも頻出の深さ優先探索(depth-first search, dfs)や幅優先探索(breadth first search、bfs)について最低限の文字数で理解できるための自分なりの解説記事を書きたくなったので執筆してみます。
なおタイトルにある通りコードの実装は全てpythonを想定しています。

グラフとは何か

dfsやbfsはその名の通り探索をするアルゴリズムですが、その探索される対象はグラフです。グラフと言っても高校までの数学に出てくるようなxy平面上にy = f(x)の挙動の絵を描いたものではありません。この世界においてグラフとはいくつかの頂点と辺を持ったデータ構造のことを指します。例えば電車の路線図であればたくさんの駅とそれらが隣り合っていることを表す線が引かれているかと思います。これはまさしくグラフであり、駅が頂点、線が辺に対応しているわけです。

競プロではよく木構造が出題されますが、木とは「閉路(ループ)を持たない」かつ「連結している(2つ以上に切断されていない)」グラフを指す言葉ですので当然dfsやbfsが有効です。しかし、それ以外のグラフ構造、時に問題を一度読んだだけではグラフに見えないかもしれない対象についてもこのアルゴリズムは効果を発揮することがあります。Atcoderで緑になるくらいまでであればグラフの問題を見たらdfs or bfsを疑ってみてもいいかもしれません。

このアルゴリズムはどのような挙動なのか

dfsは「深さ」優先です。つまり今いる頂点と繋がっている未探索の頂点がある限り進み続けます。行き止まりに当たったら初めて一つ前の頂点に戻り、別な辺を通って次の未探索な頂点を目指していきます。このようにして今いる位置から探索すべき頂点があれば進み、なければ戻る探索方法がdfsです。多くの人がRPGなどでダンジョンを探索する際にこのように探索していたのではないでしょうか。

一方bfsは「幅」優先です。すなわち現在の頂点と繋がる全ての頂点(ここでは仮にA~Cとします)をまず探索します。次にAと繋がる全ての頂点(ここではDとおきます)を探索し、Bと繋がる全ての頂点を探索し、Cと繋がる全ての頂点を探索します。ここで最初の頂点と繋がった全ての頂点を探索したのでやっとDの探索へと移ります。家系図を想像した時にある夫婦の子供達を見て、孫の代全員を見て、ひ孫の代を見て...、という順番で探索したらそれがまさにbfsです。

関数の再帰呼び出しについて

さて、dfsを考える上で欠かせないのが関数の再帰呼び出しです。

pythonではプログラムは上から順番に実行されます。その中で関数の呼び出しがあれば、一旦作業を中断した上で関数の処理を行い、終わった後にまた元の作業に戻ってきます。

def my_func(num):
    return num * 2

a = my_func(1)
print(a) # 2

なんてコードを実行した場合、一旦aにmy_funcの返り値を代入したい！という状態を保持した上でmy_funcを処理します。至って当たり前のことです。

では次のコードはどうでしょうか。

def factorial(num):
    if num < 0:
        return -1
    if num == 0:
        return 1
    
    return num * factorial(num - 1)
    
a = factorial(3)
print(a) # 6

実行の順番を追ってみましょう。
factorial(3)を呼び出してその返り値をaに代入したい -> 3は正の値なのでif文の条件には当たらず -> 返り値は3 * factorial(2)らしい -> ではfactorial(2)を呼び出し -> 2も正の値だからif文の条件には当たらず -> 返り値は2 * factorial(1)らしい -> ではfactorial(1)を呼び出し -> 1も正の値だからif文の条件には当たらず -> 返り値は1 * factorial(0)らしい -> ではfactorial(0)を呼び出し -> 2個目のif文よりnum == 0の時の返り値は1らしい -> ということはfuctorial(1)の返り値は1 * 1 ( = 1)なのか! -> ということはfuctorial(2)の返り値は2 * 1 (= 2)なのか! ->ということはfuctorial(3)の返り値は3 * 2 (=6)なのか! -> aの値は6ですよ

長いですね。コードの長さは最初とそこまでは変わらないのに一気に複雑になりました。複雑さの原因は再帰が続く限り返り値が決定しないところにあります。再帰が続く限りコンピュータはたくさんの状態を保持し続けることになります。言うなればb = (3 * (2 * (1 *(1) ) ) )といった感じです。

ところでこの、終わりが見えるまで突き進む感覚に見覚えはありませんか？これを応用するとdfsを簡単に実装することができます。dfsという関数を定義してあげて、探索できる頂点があれば次の頂点を引数に再帰呼び出しをすると次の頂点が存在する限り掘り進みます。そして最深部まできたら自動的に一つ前の関数の処理に戻って別な行き先を引数に新たな再帰呼び出しをして…と行った具合に。頂点の数だけの再帰呼び出しはfor文を使えば楽に実現可能でしょう。

なお、dfsで「最初にたどり着いた順番」と「その地点を最後に去る順番」が重要になる場面がありますが、これも再帰の工夫で解決できます。
再帰呼び出しの前の処理 -> 次の呼び出しの前の処理 -> ... -> 最深部 -> 呼び出し後の処理 -> その1つ前の再帰での関数呼び出し後の処理 -> ... -> 最初の再帰における関数呼び出し後の処理
の順に実行されるため、たどり着いた時の処理を再帰呼び出しの前に、帰る際の処理を再帰呼び出しの後に書いておくことで実現可能です。

キューというデータ構造

プログラミング初学者にとって、複数のデータを扱う際に最も身近なのは配列でしょう。配列はi番目の要素にO(1)でアクセスできるメリットがある反面、配列の途中のデータを削除しようとするとO(N)の処理時間がかかってしまいます。そのため配列以外にも様々なデータ構造が考案されています。

そのうちの一つにキュー(queue)というものがあります。これはFIFO( first in, first out)とも呼ばれているデータ構造でトンネルのようなイメージです。トンネルの入り口から車を詰めていった場合、出口から取り出す場合は最初に入った車から出てきます。i番目の要素にアクセスするにはO(i)かかるのに対し、先頭のデータを削除するのにO(1)で可能なメリットがあります。

なお、キューの仲間でスタック(stack)というものがあり、こちらはLIFO(last in, first out)などと呼ばれています。またpythonではキューとしてもスタックとしても使えるデキューというものがあり、実装の上ではこれを使うことになります。(工夫により配列を使ってもキューは実現できますが…) また、キューやスタックの良さはこれだけではありませんがbfsを考える上で使うのはこれくらいなのでご勘弁を…

bfsの考え方を少し違った表現で言い直すと現在の頂点と繋がる頂点を「探索するリスト」に加えて、「探索するリスト」を順番に潰していくと言えます。「加えて」や「順番に」がミソで、リストの最下位に追加しながら上から順番に見ていくのがbfsの考え方です。そしてこれが得意なのがまさにキューなのです。キューは入れた順番が探索の優先度合いです。最初に「子供世代」の情報を入れて次に「孫世代」の情報を入れればその順番に探索が可能というわけです。

具体的な実装例

ここではAOJのdfsとbfsの基本問題を解いていきます。

DFS

問題は上記のリンクからご覧ください。
実際にかいたコードが以下になります。

n = int(input()) # 点の総数
path = list() # path[p] : 点pから行ける点のリスト
path.append([])
for _ in range(n):
   tmp = list(map(int, input().split())) # tmp[0]から行ける点はtmp[1]個あり、それはtmp[2:]
   path.append(tmp[2:])

d = [0]*(n+1) # 最初に発見した時刻
f = [0]*(n+1) # 隣接リストを調べ終えた時刻

TIME = 0

def dfs(p, d, f):
   global TIME

   # ここに来訪時の処理を書く
   TIME += 1
   d[p] = TIME # 時刻の記録
   
   for nxt in path[p]: #繋がってる点の内
       if d[nxt] == 0: # 未探索の(発見時刻が初期値のままの場合)
           dfs(nxt, d, f) # 先に進む
   
   # ここに帰る時の処理を書く
   TIME += 1
   f[p] = TIME # 繋がってる全ての点を探索し終えたらその点でやることは終わり
   
   return     
   
for start in range(1, n+1):
   if d[start] == 0: # 未探索の点があれば
       dfs(start, d, f) # dfs開始

for i in range(1, n+1):
   print(i, d[i], f[i])

この問題では時刻を記録する必要があったのでdとfという配列を用意していますが、ただ探索するだけならfに相当する配列は不要ですし、閉路を持たないことが保証されていれば(例えば木構造なら)来訪を記録する配列は必要ないことも多々あります。

BFS

こちらも問題はリンク先を確認ください。実装は以下の通りです。

from collections import deque
n = int(input())
ukv = [list(map(int, input().split())) for _ in range(n)]

distance = [-1]*(n+1) # 頂点1からの距離を記録

q = deque() # キューを宣言して
q.append(1) # 探索開始地点を追加

# bfs
distance[1] = 0 # 距離の初期化
while len(q) > 0:
   now = q.popleft() # 今探索する点の情報の取得(popした時点でqから値は消えている)

   for i in range(ukv[now - 1][1]): # ukv[now - 1][1] : nowから出てるパスの数
       node = ukv[now - 1][2+i] # node = 今注目している頂点
       if distance[node] == -1: # nodeが未探索ならば
           q.append(node) # q(探索リスト)に追加する
           distance[node] = distance[now] + 1 # bfsでは一つ前にどの頂点にいたかの情報がqからpopした時点ではわからないのでqに追加する段階で更新しておくと良い
   

for i in range(1, n+1):
   print(i, distance[i])

bfsは開始地点から距離が近い順番に探索していく方式ですのでdistance[node]を更新する際にmin関数などを噛ませなくて良いのが特徴ですね。またdistanceを更新する場所も重要で、for文の外で更新をしようとしても前にどの頂点にいたのかを簡単には調べられません。そのためqに追加する時に必要な処理をしておくと楽に実装できます。