木構造の幅優先探索・深さ優先探索コピペ用
tl;dr
AtCoderで木構造の幅優先探索、深さ優先探索が出た場合のコピペ用ソースコードを記す。
解けなかった問題
これ
幅優先探索+色探索を実装していたらTLEしか出ませんでした…。
幅優先探索
木構造があったとき、こんな感じでとけます。
木
pythonのdefaultdictで表現します。
# Atcoderでよく与えられる形 T_list = [[1,2],[2,3],[3,4],[3,5],[2,7],[7,8],[7,9]] # defaultdictでtreeを表現 from collections import defaultdict as ddict tree = ddict(list) for t in T_list: tree[t[0]].append(t[1]) print(tree)
こんな感じで表現されます。
defaultdict(<class 'list'>, {1: [2], 2: [3, 7], 3: [4, 5], 7: [8, 9]})
幅優先探索関数
pythonのdequeでやります。
from collections import deque def bfs(T,i): Q = deque([i]) print(i) while Q: q = Q.popleft() for c in T[q]: print(c) Q.append(c)
使い方は第一引数に木を、第二引数にどこから探索するかを入れます。
# 1から探索 bfs(tree,1)
結果は下記の通り
1 2 3 7 4 5 8 9
ちゃんと階層が浅い順に探索されています。
深さ優先探索
今回の問題ではなかったのですが、ついでに深さ優先探索も。
深さ優先探索関数
再帰を利用して実装。
def dfs(T,i,init=True): if init==True: print(i) Q = deque([i]) q = Q.popleft() for c in T[q]: print(c) dfs(T,c,False)
使い方は幅優先探索と一緒だけれども、
探索開始地点の出力のために、
第三引数を用意します。(初回のみTrueで出力)
from collections import deque def dfs(T,i,init=True): if init==True: print(i) Q = deque([i]) q = Q.popleft() for c in T[q]: print(c) dfs(T,c,False)
こんな感じで使えます。
dfs(tree,1)
結果もちゃんと深さ優先で出ています。
1 2 3 4 5 7 8 9