Pythonで木構造の問題を解いてみた

こんにちは，今回は前回扱った木構造の復習を兼ねて木を扱った問題と木の仕組みを利用した問題二つを解いていきたいと思います．
まだ，前回記事をみていない方，木構造を知らない方は一つ前の記事を読んでみてください．

目次

木構造の練習問題1(木の巡回)

今回の問題はALDS1_7_C: Tree Walk です．
問題の内容としては以下の二分木の探索方法を三つ実装するというものです．

1. 根節点，左部分木，右部分木の順で節点の番号を出力する．(先行順巡回 ： Preorder Tree Walk)
2. 左部分木，根節点，右部分木の順で節点の番号を出力する．(中間順巡回 ： Inorder Tree Walk)
3. 左部分木，右部分木，根節点の順で節点の番号を出力する．(後行順巡回 : Postorder Tree Walk)

入力は最初の行に節点の個数nが与えられ，続くん行には各節点の情報(id, left, right)が与えられます．
※子供を持たない場合left, rightは-1が与えられます．

出力は上記の三つの探索方法の名称と巡回した順番に節点番号を出力するというものです．
では，練習なので早速取り掛かっていきましょう．

実装コード

class Node:
    def __init__(self):
        self.parent = -1

def make_orderlist(node_id, answer_dic):
    preorder(node_id, answer_dic)
    inorder(node_id, answer_dic)
    postorder(node_id,answer_dic)

def preorder(node_id, answer_dic):
    if node_id == -1:
        return
    answer_dic["Preorder"].append(str(node_id))
    preorder(Tree[node_id].left, answer_dic)
    preorder(Tree[node_id].right, answer_dic)
    
def inorder(node_id, answer_dic):
    if node_id == -1:
        return
    inorder(Tree[node_id].left, answer_dic)
    answer_dic["Inorder"].append(str(node_id))
    inorder(Tree[node_id].right, answer_dic)

def postorder(node_id, answer_dic):
    if node_id == -1:
        return
    postorder(Tree[node_id].left, answer_dic)
    postorder(Tree[node_id].right, answer_dic)
    answer_dic["Postorder"].append(str(node_id))
    

N = int(input())
Tree = [Node() for _ in range(N)]

#make_tree
for _ in range(N):
    #id, left, right
    tree_info = list(map(int, input().split()))
    node_id = tree_info[0]
    left = tree_info[1]
    right = tree_info[2]
    #自身の子をセット
    Tree[node_id].left = left
    Tree[node_id].right = right
    #子からみた親をセット
    if left != -1:
        Tree[left].parent = node_id
    if right != -1:
        Tree[right].parent = node_id

        
#search_root
root_id =[i for i,t in enumerate(Tree) if t.parent == -1][0]
answer_dic = {"Preorder":[], "Inorder":[], "Postorder":[]}
make_orderlist(root_id, answer_dic)

#answer_output
for i in answer_dic:    
    print(i)
    #出力の形式上先に空白を一つ入れる必要があるらしい
    print(" "+" ".join(answer_dic[i]))

コードの解説

前回と同じコードを最初の部分にはそのまま活用しています．
今回扱っている木構造では兄弟の情報や深さなどは必要なく，単純に自分の親と左右の子供の状態を知っていれば大丈夫です．
また，基本的にはrootは最初の入力された節点だと思うので必要に応じて使い分けてください．
中身に関してはmake_orderlistの部分が新しい部分にはなってきます．
しかし，ここの部分に関しても今までやってきた再帰のやり方を覚えていれば難なくできると思います．

練習問題1の応用(木の巡回の応用)

問題内容

ある二分木に対して，それぞれ先行順巡回と中間順巡回を行なって得られる節点番号の列が与えられるので，その二分木の後行順巡回で得られる節点の列を出力するプログラムを作成するというものです．
入力としては一行目に節点の数nが与えられ，二行目に先行順巡回で得られる節点番号の列，三行目に中間順巡回で得られる節点番号の列が与えられます． 出力は後行順巡回で得られる節点番号の列を出力します．
問題は理解できますが，実装方法が若干複雑そうですね…

問題の解説

問題を順番に考えていきたいと思います． 先行順巡回では根→左部分木→右部分木の順に巡回していきます．
一方，中間順巡回では左部分木→根→右部分木の順に巡回していきます．
それぞれの巡回で生まれる違いを利用していきたいと思います．

まず，上記の図のような木構造が与えられていた時
先行順巡回 = [1,2,3,4,5,6,7,8,9]
中間順巡回 = [3,2,5,4,6,1,8,7,9]
というような列が与えられます．
この時，先行順巡回の先頭の値は根であるので中間順巡回を
左部分木 = [3,2,5,4,6]
右部分木 = [8,7,9]
のように分けることができます.
次に先行順巡回の次の値は"2"なので左部分木をさらに
左部分木 = [3]
右部分木 = [5,4,6]
のようにわけることができます．
この操作を繰り返すことで木構造がどのような形であったか再構築することができます．
この再構築をしていく過程で後行順巡回の巡回方法に沿って答えのリストを作っていけば問題を解くことができそうですね.
では実際に実装していきたいと思います.

実装コード

def reconstruct(l, r):
    if l >= r:
        return
    root = pre_list.pop(0)
    center = in_list.index(root)
    reconstruct(l, center) # 左部分木
    reconstruct(center+1, r) #右部分木
    ans_list.append(str(root))
    
N = int(input())
pre_list = list(map(int, input().split()))
in_list = list(map(int, input().split()))

#answer_output
ans_list = []
reconstruct(0, N)
print(" ".join(ans_list))

コードの解説

コードに落とし込むと解説するほどのものではなくなっていますね.
先行順巡回の節点を引っ張り，それに対応する中間順巡回の場所をまずは探します．
その中間順巡回の場所が左部分木と右部分木にわける位置となっています．
同様の操作を分けられた左部分木と右部分木でも行うよう再帰させています．
左部分木が再帰できなくなった点が最初の後行順巡回の値になるのでそのタイミングでリストに値を追加されるようにしています.

最後に

今回は前回扱った木構造の問題を復習しながら扱いました．また，扱った木の巡回の応用問題も解いていきました．
この木の巡回の応用は二分探索にもつながる考え方なので二分探索の仕組みを知りたい人にぜひ解いてもらいたいです．

解いてほしい問題や解き方に関するコメントなどなんでもいいのでありましたらコメントよろしくお願いします!
スターも励みになるので気が向いたら押してください！