AOJ Lesson ALDS1 7_D "Reconstruction of a Tree" をPythonで解いたみた #Python

AOJのLessonのReconstruction of a Treeを解いた際の備忘録です。
「AOJにある解答例は難しすぎる、、ネットでもわかりやすい解説もない、、」
という状況だったので、素人である自分に合わせた可読性の高い(?)コードを残しておきます。

問題

AOJ Lesson アルゴリズムとデータ構造 Reconstruction of a Tree

ある二分木に対して、それぞれ先行順巡回 (preorder tree walk) と中間順巡回 (inorder tree walk) を行って得られる節点の列が与えられるので、その二分木の後行順巡回 (postorder tree walk) で得られる節点の列を出力するプログラムを作成してください。

入力

１行目に二分木の節点の数nが与えられます。
２行目に先行順巡回で得られる節点の番号の列が空白区切りで与えられます。
３行目に中間順巡回で得られる節点の番号の列が空白区切りで与えられます。

入力例

5
1 2 3 4 5
3 2 4 1 5

出力例

3 4 2 5 1

解法

自分の実力は通用せず(涙)
以下の記事がわかりやすかったです。

大事な部分だけ取り出すと、

ある部分木について、preorderから、節点を見つける
見つけた節点でinorderを分割し、左部分木、右部分木を作る（実際はinorderの列）
2で作った左右それぞれの部分木に対応するpreorderの列を生成する。
2,3で作った左右それぞれのpreorderとinorderの列に対して、1を再帰的に実行する

という流れのコードを組めれば良いということです。

実装

上で紹介した記事を元にpythonで実装してみました

from sys import stdin
from typing import List


class Node(object):
    def __init__(self):
        self.left_child = None
        self.right_child = None
        self.parent = -1


n = int(stdin.readline())
*preorder, = map(int, stdin.readline().split())
*inorder, = map(int, stdin.readline().split())

tree = [Node() for _ in range(n+1)]


def reconstruction(inorder: List[int], preorder: List[int], parent=-1):
    if len(preorder) == 0:
        return None
    # ある部分木について、preorderから、節点を見つける
    root = preorder[0]
    current_node = tree[root]
    # nodeのparentの登録
    current_node.parent = parent

    if len(preorder) == 1:
        return root

    root_index_in_inorder = inorder.index(root)
    # 見つけた節点でinorderを分割し、左部分木、右部分木を作る
    left_inorder = inorder[:root_index_in_inorder]
    right_inorder = inorder[root_index_in_inorder+1:]

    # 上で作った左右それぞれの部分木に対応するpreorderの列を生成する。
    left_preorder = [
        i for i in preorder if i in left_inorder]
    right_preorder = [
        i for i in preorder if i in right_inorder]

    # 作った左右それぞれのpreorderとinorderの列に対して、1を再帰的に実行する
    current_node.left_child = reconstruction(
        left_inorder, left_preorder, root)
    current_node.right_child = reconstruction(
        right_inorder, right_preorder, root)

    # rootを返す(この値が呼び出し元のnodeのchildとなる)
    return root


root_id = preorder[0]
reconstruction(inorder, preorder)

# 出力部分
result = []
def postorderWalk(root_id: int):
    if root_id is not None:
        node = tree[root_id]
        postorderWalk(node.left_child)
        postorderWalk(node.right_child)
        result.append(root_id)


postorderWalk(root_id)
print(*result)