木構造の幅優先探索・深さ優先探索コピペ用

tl;dr

AtCoder木構造幅優先探索深さ優先探索が出た場合のコピペ用ソースコードを記す。

解けなかった問題

これ

atcoder.jp

幅優先探索+色探索を実装していたら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

Pythonで学ぶアルゴリズムとデータ構造 (データサイエンス入門シリーズ)

Pythonで学ぶアルゴリズムとデータ構造 (データサイエンス入門シリーズ)

  • 作者:辻 真吾
  • 出版社/メーカー: 講談社
  • 発売日: 2019/11/28
  • メディア: 単行本(ソフトカバー)