Pythonで二分木の先行順走査（プレオーダートラバーサル）を実装する方法

Pythonでの二分木の先行順走査（プレオーダートラバーサル）とは

二分木が与えられたとき、その木を先行順走査（プレオーダートラバーサル）で巡回した結果を返すことを考えます。先行順走査とは、「根のノード → 左部分木 → 右部分木」の順序でノードを訪問する木構造の基本的な走査手法です。

例えば、次のような二分木があるとします。

この木に対する先行順走査の結果は [3, 9, 20, 15, 7] となります。

アルゴリズムの手順

ここでは、再帰を使わずにスタックを利用した反復的なアプローチで問題を解きます。手順は以下の通りです。

• 結果を格納するための空リスト res と、スタックとして使用する空リスト st を用意します。
• node をルートに設定します。
• node が存在するか、st が空でない限り、以下の処理を繰り返します。
  • node が null でない間、次の操作を行います。
    • node の値を res に追加し、node 自身を st にプッシュしてから、node を左の子へ移動させます。
  • st の末尾の要素を取り出して temp とします。
  • temp に右の子が存在する場合は、node をその右の子に設定します。
• 最後に res を返します。

実装例

それでは、実際のPythonコードを見てみましょう。

サンプルコード

class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.data = data
        self.left = left
        self.right = right
def insert(temp,data):
    que = []
    que.append(temp)
    while (len(que)):
        temp = que[0]
        que.pop(0)
        if (not temp.left):
            temp.left = TreeNode(data)
            break
        else:
            que.append(temp.left)
        if (not temp.right):
            temp.right = TreeNode(data)
            break
        else:
            que.append(temp.right)
def make_tree(elements):
    Tree = TreeNode(elements[0])
    for element in elements[1:]:
    insert(Tree, element)
    return Tree
class Solution(object):
    def preorderTraversal(self, root):
        res = []
        st = []
        node = root
        while node or st:
            while node:
                if node.data != None:
                    res.append(node.data)
                st.append(node)
                node = node.left
            temp = st[-1]
            st.pop()
            if temp.right:
                node = temp.right
        return res
ob1 = Solution()
head = make_tree([3,9,20,None,None,15,7])
print(ob1.preorderTraversal(head))

入力

[3,9,20,null,null,15,7]

出力

[3, 9, 20, 15, 7]

まとめ

このように、スタックを活用することで再帰呼び出しに頼らずに先行順走査を実装できます。計算量はノード数を n とすると、時間計算量・空間計算量ともに O(n) となり、すべてのノードを一度ずつ訪問する効率的なアルゴリズムです。再帰版と併せて理解しておくと、木構造の走査に関する応用問題にも対応しやすくなります。