バイナリツリーのpre/in/post-order DFS（再帰あり/なし）

概要

DFSの練習で、Pythonでバイナリツリーの探索を実装しました。

再帰あり/なしのそれぞれで、preorder, inorder, postorderの計6パターンの実装です。 再帰なしのpostorderは下記GeeksforGeeksの記事を参考にしています。

www.geeksforgeeks.org

コード全文はgithubに上げています。 github.com

バイナリツリー生成

バイナリツリーは1以上の深さ指定で自動生成しています。 クラスを定義しようかとも思いましたが、とりあえずサクッとリストで表現しました。

def make_binarytree(depth):
    tree = []
    for i in range(1, 2**depth-2, 2):
        tree.append([i, i+1])
    return tree

# make_binarytree(3)
# [[1, 2], [3, 4], [5, 6]]

DFS

preorder, inorder, postorderのそれぞれで再帰あり/なしのDFSを実装しています。 再帰ありは比較的きれいに書けたかなと思いますが、 スタックを使った再帰なしのinorderは我ながら汚いです。 再帰なしpostorderはいい方法が思いつかなかったので、上記GeeksforGeeksの記事を参考にしました。

inorderは、探索用とは別に表示順管理用のスタックを用意しています。

postorderは、各ノードで探索スタック最上位のノードを別スタックに積んでいくと、 最上位から表示したい順に積まれていきます。 右側から探索していくのがミソなようです。

class DFS():
    def __init__(self, tree):
        self.tree = tree
    
    def preorder_wi_recursive(self):
        def _inner(parent):
            print(parent, '', end='')
            if parent > len(self.tree) - 1: return
            _inner(self.tree[parent][0])
            _inner(self.tree[parent][1])
        _inner(0)
        print()

    def inorder_wi_recursive(self):
        def _inner(parent):
            if parent > len(self.tree) - 1:
                print(parent, '', end='')
                return
            _inner(self.tree[parent][0])
            print(parent, '', end='')
            _inner(self.tree[parent][1])
        _inner(0)
        print()

    def postorder_wi_recursive(self):
        def _inner(parent):
            if parent > len(self.tree) - 1:
                print(parent, '', end='')
                return
            _inner(self.tree[parent][0])
            _inner(self.tree[parent][1])
            print(parent, '', end='')
        _inner(0)
        print()
    
    def preorder_wo_recursive(self):
        stack = deque()
        stack.append(0)
        while len(stack) is not 0:
            parent = stack[-1]
            stack.pop()
            print(parent, '', end='')
            if parent < len(self.tree):
                stack.append(self.tree[parent][1])
                stack.append(self.tree[parent][0])
        print()

    def inorder_wo_recursive(self):
        stack = deque()
        prev = deque()
        stack.append(0)
        while len(stack) is not 0:
            parent = stack[-1]
            stack.pop()
            if parent < len(self.tree):
                prev.append(parent)
                stack.append(self.tree[parent][1])
                stack.append(self.tree[parent][0])
            else:
                print(parent, '', end='')
                if len(prev):
                    print(prev.pop(), '', end='')
        print()

    def postorder_wo_recursive(self):
        stack = deque()
        prev = deque()
        stack.append(0)
        while len(stack) is not 0:
            prev.append(stack[-1])
            parent = stack.pop()
            if parent < len(self.tree):
                stack.append(self.tree[parent][0])
                stack.append(self.tree[parent][1])
        prev = list(prev)
        prev.reverse()
        print(' '.join(list(map(str, prev))))

if __name__ == '__main__':
    depth = int(input())
    dfs = DFS(make_binarytree(depth))
    dfs.preorder_wi_recursive()
    dfs.inorder_wi_recursive()
    dfs.postorder_wi_recursive()
    dfs.preorder_wo_recursive()
    dfs.inorder_wo_recursive()
    dfs.postorder_wo_recursive()

実行例

>python3 binarytree.py 
3
0 1 3 4 2 5 6 
3 1 4 0 5 2 6 
3 4 1 5 6 2 0 
0 1 3 4 2 5 6 
3 1 4 0 5 2 6 
3 4 1 5 6 2 0