今更感ありますが、、二分木構造に関するDFSで、1. 行きがけ、2. 帰りがけ, 3. 全記録、のPythonコードを紹介します。本稿では下記の木を例に使用します。
幅優先探索(BFS)はオレンジ色の順番のように、深さが浅いノードを優先して探索を行います。
ここでは走査した順でノード(頂点)番号を記録することが要求されているとします。基本的にはqueue(first in, first out)データ構造を用いて実装します。
配列で二分木情報が与えられている場合
root = [6,2,8,0,4,7,9,None,None,3,5]
index iの左の子は2i+1, 右の子は2i+2, 根はindex 0
幅優先探索
import collections queue = collections.dequeu() def bfs(i, root): while queue: yield root[i] if i < (len(root) >> 1): left = (i << 1) + 1 if root[left] is not None: queue.append(left) right = (i << 1) + 2 if root[right] is not None: queue.append(right) print(list(dfs(0, root))) >>> [6, 2, 0, 4, 3, 5, 8, 7, 9]
最良優先探索
最良優先探索として数字が小さいノードを優先して探索するコードは下記になります。優先度つきキュー(priority queue)としてheapqを用いています。これは管理するキューの中で最大の要素を効率よく取り出す操作をサポートします。
import heapq def score(i): return root[i] def bfs(root): priority_queue = [(score(0), 0)] heapq.heapify(priority_queue) while priority_queue: score_value, i = heapq.heappop(priority_queue) yield root[i] if i < (len(root) >> 1): left = (i << 1) + 1 if root[left] is not None: heapq.heappush(priority_queue, (score(left), left)) right = (i << 1) + 2 if root[right] is not None: heapq.heappush(priority_queue, (score(right), right)) list(bfs(root)) >>> [6, 2, 0, 4, 3, 5, 8, 7, 9]
Nodeクラスで情報が与えられている場合
次のTreeNodeクラスのrootが与えられている場合。属性のleft, rightには左側、右側の子に相当するTreeNodeが格納されており、もし子がいない場合はNoneが格納されています。
class TreeNode: def __init__(self, x): self.val = x self.left = None self.right = None
上記の二分木の場合。
root = [6,2,8,0,4,7,9,None,None,3,5] TreeNodes = list() for i, x in enumerate(root): if x is None: TreeNodes.append(None) else: node = TreeNode(x) TreeNodes.append(node) if i > 0: parent = TreeNodes[(i + 1) // 2 - 1] if i % 2: parent.left = node else: parent.right = node root_node = TreeNodes[0]
幅優先探索
def bfs(root): queue = collections.deque() queue.append(root) while queue: node = queue.popleft() yield node if node.left is not None: queue.append(node.left) if node.right is not None: queue.append(node.right) print([node.val for node in bfs(root_node)]) >>> [6, 2, 8, 0, 4, 7, 9, 3, 5]
最良優先探索
import heapq def score(node): return node.val def bfs(root): priority_queue = [(score(root), root)] heapq.heapify(priority_queue) while priority_queue: score_value, node = heapq.heappop(priority_queue) yield node if node.left is not None: heapq.heappush(priority_queue, (score(node.left), node.left)) if node.right is not None: heapq.heappush(priority_queue, (score(node.right), node.right)) list(bfs(root)) >>> [6, 2, 0, 4, 3, 5, 8, 7, 9]
関連
深さ優先探索(DFS)