サブロウ丸

サブロウ丸

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

(DFS) 深さ優先探索; Python

二分木構造に関するDFSで、1. 行きがけ、2. 帰りがけ、 3. 通りがけ、 4. 全記録、のPythonコードを紹介します。本稿では下記の木を例に使用します。

深さ優先探索(DFS)はご存知のようにオレンジ色の順番のように、深く(子孫側に)進む方向を優先して二分木を走査します。

ここでは走査した順でノード(頂点)番号を記録することが要求されているとします。そのときに、行きがけの順(preorder traversal)とは初めてそのノードに到達した順にノード番号を並べることで、帰りがけの順(postorder traversal)は最後にそのノードに到着した順、通りがけ(inorder traversal)は左から右にノードを並べた順、全記録はとにかくそのノードを通ったらノード番号を記録することにします。

行きがけ, 帰りがけ, 通りがけの順では同一のノード番号が一回限り記録されるのに対し、全記録は最大で3回記録されます。(ちなみに全記録はこの筆者が適当につけたネーミングです)

配列で二分木情報が与えられている場合

root = [6,2,8,0,4,7,9,None,None,3,5]

index iの左の子は2i+1, 右の子は2i+2, 根はindex 0

行きがけ (preorder traversal)

def dfs(i, root):
    yield root[i]
    if i < (len(root) >> 1):
        left = (i << 1) + 1
        if root[left] is not None:
            yield from dfs(left, root)
        right = (i << 1) + 2
        if root[right] is not None:
            yield from dfs(right, root)

print(list(dfs(0, root)))
>>> [6, 2, 0, 4, 3, 5, 8, 7, 9]

帰りがけ (postorder traversal)

def dfs(i, root):
    if i < (len(root) >> 1):
        left = (i << 1) + 1
        if root[left] is not None:
            yield from dfs(left, root)
        right = (i << 1) + 2
        if root[right] is not None:
            yield from dfs(right, root)
    yield root[i]

print(list(dfs(0, root)))
>>> [0, 3, 5, 4, 2, 7, 9, 8, 6]

通りがけ (inorder traversal)

def dfs(i, root):
    if i < (len(root) >> 1):
        if root[left] is not None:
            yield from dfs(left, root)
        right = (i << 1) + 2
        if root[right] is not None:
            yield from dfs(right, root)
    yield root[i]

print(list(dfs(0, root)))
>>> [0, 2, 3, 4, 5, 6, 7, 8, 9]

全記録

def dfs(i, root):
    yield root[i]
    if i < (len(root) >> 1):
        left = (i << 1) + 1
        if root[left] is not None:
            yield from dfs(left, root)
            yield root[i]
        right = (i << 1) + 2
        if root[right] is not None:
            yield from dfs(right, root)
            yield root[i]

print(list(dfs(0, root)))
>>> [6, 2, 0, 2, 4, 3, 4, 5, 4, 2, 6, 8, 7, 8, 9, 8, 6]



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]

行きがけ (preorder traversal)

def dfs(node):
    yield node
    if node.left is not None:
        yield from dfs(node.left)
    if node.right is not None:
        yield from dfs(node.right)

print([node.val for node in dfs(root_node)])
>>> [6, 2, 0, 4, 3, 5, 8, 7, 9]

帰りがけ (postorder traversal)

def dfs(node):
    if node.left is not None:
        yield from dfs(node.left)
    if node.right is not None:
        yield from dfs(node.right)
    yield node

print([node.val for node in dfs(root_node)])
>>> [0, 3, 5, 4, 2, 7, 9, 8, 6]

通りがけ (inorder traversal)

def dfs(node):
    if node.left is not None:
        yield from dfs(node.left)
    yield node
    if node.right is not None:
        yield from dfs(node.right)
    yield node

print([node.val for node in dfs(root_node)])
>>> [0, 2, 3, 4, 5, 6, 7, 8, 9]

全記録

def dfs(node):
    yield node
    if node.left is not None:
        yield from dfs(node.left)
        yield node
    if node.right is not None:
        yield from dfs(node.right)
        yield node

print([node.val for node in dfs(root_node)])
>>> [6, 2, 0, 2, 4, 3, 4, 5, 4, 2, 6, 8, 7, 8, 9, 8, 6]

P.S.

yield from の使用法を紹介するのが実はサブテーマ。

関連