サブロウ丸

サブロウ丸

主にプログラミングと数学

(BFS) 幅優先探索; python

今更感ありますが、、二分木構造に関する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)