Pythonで先行順(プレオーダー)と後行順(ポストオーダー)から二分木を構築する方法

はじめに

二分木の復元問題は、コーディング面接でも頻出のテーマです。本記事では、先行順トラバーサル(プレオーダー)と後行順トラバーサル(ポストオーダー)という2つの走査結果が与えられたときに、元の二分木を再構築する方法をPythonで解説します。

たとえば、先行順が [1,2,4,5,3,6,7]、後行順が [4,5,2,6,7,3,1] の場合、次のような二分木が得られます。

一意な復元に関する注意点

まず押さえておきたいのは、先行順と後行順の組み合わせだけでは、すべての内部ノードが2つの子を持つ場合に限り木が一意に決定されるという点です。子を1つしか持たないノードが存在すると、その子が左側なのか右側なのかを判別できません。本記事の実装では、慣例に従って左側を優先して配置する方針を取ります。

アルゴリズムの手順

スタックを活用することで、この問題を線形時間で解くことができます。具体的な手順は以下の通りです。

• ans: pre[0] を値とするノードを作成する。空のスタックを用意し、ans をプッシュする。
• i := 1(先行順の読み取り位置)、j := 0(後行順の読み取り位置)で初期化する。
• i < len(pre) かつ j < len(post) の間、以下を繰り返す。
  • スタックの先頭ノードの値が post[j] と等しい場合:j を1増やし、スタックからポップして、次の反復へ進む。
  • pre[i] を値とする新しいノード node を作成する。
  • スタック先頭ノードの左の子が空であれば node を左の子に設定し、そうでなければ右の子に設定する。
  • node をスタックにプッシュする。
  • i を1増やす。
• 最後に ans を返す。

なぜスタックで解けるのか

先行順は「根 → 左 → 右」、後行順は「左 → 右 → 根」の順でノードを訪問します。スタックの先頭にあるノードの値が後行順の現在位置 post[j] と一致したとき、そのノードを根とする部分木はすでに完成したことを意味します。この性質を利用すれば、「どのノードに新しい子を追加すべきか」をスタックの先頭だけで判断できるため、効率的に木を組み立てられます。

Pythonでの実装例

以下が実際のコードです。結果をレベル順(幅優先)で出力するための補助関数も含めています。

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

def height(root):
    if root is None:
        return 0
    else:
        # 左右の部分木の高さをそれぞれ計算
        l_height = height(root.left)
        r_height = height(root.right)
        # 大きい方に1を加えて返す
        if l_height > r_height:
            return l_height + 1
        else:
            return r_height + 1

def print_given_level(root, level):
    if root is None:
        return
    if level == 1:
        print(root.data, end=',')
    elif level > 1:
        print_given_level(root.left, level - 1)
        print_given_level(root.right, level - 1)

def level_order(root):
    print('[', end='')
    h = height(root)
    for i in range(1, h + 1):
        print_given_level(root, i)
    print(']')

class Solution(object):
    def constructFromPrePost(self, pre, post):
        """
        :type pre: List[int]
        :type post: List[int]
        :rtype: TreeNode
        """
        ans = TreeNode(pre[0])
        stack = [ans]
        i = 1
        j = 0
        while i < len(pre) and j < len(post):
            if stack[-1].data == post[j]:
                j += 1
                stack.pop(-1)
                continue
            node = TreeNode(pre[i])
            if not stack[-1].left:
                stack[-1].left = node
            else:
                stack[-1].right = node
            stack.append(node)
            i += 1
        return ans

ob = Solution()
pre = [1,2,4,5,3,6,7]
post = [4,5,2,6,7,3,1]
tree = ob.constructFromPrePost(pre, post)
level_order(tree)

入力

pre = [1,2,4,5,3,6,7]
post = [4,5,2,6,7,3,1]

出力

[1,2,3,4,5,6,7]

末尾にカンマが付いているのは、出力時に end=',' を指定しているためです。

計算量

• 時間計算量: O(n) — 各ノードを一度ずつ処理します。
• 空間計算量: O(n) — スタックおよび再帰呼び出しのための領域が必要です。

まとめ

スタックを使うことで、先行順と後行順の2つの走査列から二分木を効率的に再構築できます。ポイントは「スタック先頭の値が後行順の現在位置と一致したら、その部分木は完成」という判定です。ただし、子を1つしか持たないノードが存在する場合は左右の区別ができないため、この手法はすべての内部ノードが2つの子を持つような木に対して特に有効であることを覚えておきましょう。