スズメの本棚

自分の知識の定着のために本や実装したことをまとめていきます.

Pythonで木構造を実装してみた ~二分木と根付き木~

こんにちは,今回は競技プログラミングでよく見る木構造について解説していきたいと思います.

目次

木構造は階層的な構造を表すのに適したデータ構造で効率的なアルゴリズムやデータ構造を実装するための基本となる考え方です.
今回は木をどのように表現するかや木構造を利用した基本的なアルゴリズムの実装をしていきたいと思います.

まず木構造とは節点(node)と節点同士を結ぶ辺(edge)で表されるデータ構造です.それを踏まえた上で根付き木と二分木を紹介していきたいと思います.

根付き木

根付き木は木の中でも根(root)と呼ばれる他の節点と区別された一つの節点を持つ木の構造のことを言います.
根付き木の節点には親子関係があります.rootを基準にそこからつながる節点は子供となり,同じ深さにある節点は兄弟と呼びます.
また,子を持たない節点を葉(leaf)や外部節点と呼びます.
図を見れば簡単に理解できると思うので載せておきます.

f:id:tech_shelf:20200803034542p:plain
木 (数学) - Wikipedia

では,実際に問題を解いていきたいと思います.

根付き木の問題

今回扱う問題はALDS1_7_A:Rooted Treesになります.
問題の内容としては与えられた根付き木の各節点について

  • 節点番号
  • 節点の種類(根, 内部ノード, 葉)
  • 親の節点番号
  • 子のリスト
  • 深さ

の情報を出力するプログラムを生成するというものです.
入力は節点の個数nとそれに続く行で節点の番号(id), 次数(k), 1~k番目までの子の節点番号(c1,…,ck)の情報が与えられます.
制約

  • 1 ≦ n ≦ 100,000
  • 節点の深さは20以下
  • 任意の二つの節点間に必ず経路が存在する

となっています.

根付き木の実装コード

では,早速実装していきたいと思います.

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

def cal_depth(node_id, d = 0):
    Tree[node_id].depth = d
    for child in Tree[node_id].children:
        cal_depth(child, d+1)
    
N = int(input())
Tree = [Node() for _ in range(N)]

#make_tree
for _ in range(N):
    #id, 子供の数k, c_0~c_k
    tree_info = list(map(int, input().split()))
    node_id = tree_info[0]
    k = tree_info[1]
    if k > 0:
        children = tree_info[2:]
        Tree[node_id].children = children
        Tree[node_id].type = "internal node"
    else:
        Tree[node_id].type = "leaf"
        
    for child in Tree[node_id].children:
        Tree[child].parent = node_id

#search_root
root_id =[i for i,t in enumerate(Tree) if t.parent == -1][0]
Tree[root_id].type = "root"
cal_depth(root_id)


#answer_output
for i, t in enumerate(Tree):
    print("node {}: parent = {}, depth = {}, {}, {}".format(i, t.parent, t.depth, t.type, t.children))

コードの解説

まず,今回扱う問題ではrootは"-1"にするように,また,節点に子がいなければ"[]"を出力するように求められているので初期化して最初の呼び出しの時にはその状態になるようにしています.

次に,深さの計算の説明をしていきます.
深さに関しては上記のプログラムの特性上すでにrootの節点が分かっている状態になっています.
そのため,root節点から上から順に辿っていくことで各節点の深さを求めることができます.
各ノードには自身の持つ子の状態が保存されているので,各子供に対して再帰的にプログラムを呼び出すことで深さを求めています

最後に,メインとなっている木構造を作っている部分に関してです.
ここでは,まずはじめにN個のNodeを作成します.
次に,入力を受け取りながら子供の情報,ノードのタイプを作っていきます.
子供に関してはkが1以上なら入力情報に合わせて子供のリストを作成,kが0なら子供がいないのでそのままとしています.
ノードのタイプに関しては一つでも子がいれば"internal node"になるのでNodeに内部ノードであると情報を追加します.
一方で,もし子供がいなければその節点は葉であることが確定しているのでNodeに対して"leaf"という情報を追加しています.
これらだけではどれがrootかというのがわかっていないので最後にすべてのノードの情報を確認し,親が"-1"のままのノードを探し,rootとしています.

二分木

二分木は根付き木の中でも全ての節点についてその子の数が2以下である木構造のことを言います.
また,二分木では節点が持つ子を左の子と右の子に区別されます.つまり,子が一つの場合にも左の子にするか右の子にするか区別する必要があります.
こちらに関しても図で確認した方がわかりやすいので載せておきます.

f:id:tech_shelf:20200803040459p:plain
二分木 - Wikipedia

二分木の問題

二分木に関して扱っている問題ALDS1_7_B: Binary_Treeを解いていきたいと思います.
問題の内容としては
与えられた二分木Tの各節点uについて

  • uの節点番号
  • uの親
  • uの兄弟
  • uの子の数
  • uの深さ
  • uの高さ
  • uの節点の種類

を出力するというものです.
入力は id, right, leftが与えられ,もし子を持たない場合left(right)-1と表されます.
問題の作り自体は先ほどの根付き木に似ているので上記のプログラムを書き換えていくこと解けるのではないかと思います.

二分木の実装コード

では,早速実装に取り掛かりたいと思います.

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

def cal_depth(node_id, d = 0):
    Tree[node_id].depth = d
    if Tree[node_id].left !=-1:
        cal_depth(Tree[node_id].left, d+1)
    if Tree[node_id].right !=-1:
        cal_depth(Tree[node_id].right, d+1)
        
def cal_height(node_id):
    left_h = 0
    right_h = 0
    if Tree[node_id].left != -1:
        left_h = cal_height(Tree[node_id].left) + 1
    if Tree[node_id].right != -1:
        right_h = cal_height(Tree[node_id].right) + 1
    Tree[node_id].height = max(left_h, right_h)
    return max(left_h, right_h)
     
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
    #ノードのタイプをセット
    if left == -1 and right == -1:
        Tree[node_id].type = "leaf"
    elif left == -1 or right == -1:
        Tree[node_id].degree = 1
        Tree[node_id].type = "internal node"
    else:
        #兄弟をセット
        Tree[right].sibling = left
        Tree[left].sibling = right
        Tree[node_id].degree = 2
        Tree[node_id].type = "internal node"
        
#search_root
root_id =[i for i,t in enumerate(Tree) if t.parent == -1][0]
Tree[root_id].type = "root"
cal_depth(root_id)
cal_height(root_id)

#answer_output
for i, t in enumerate(Tree):
    print("node {}: parent = {}, sibling = {}, degree = {}, depth = {}, height = {}, {}".format(i, t.parent, t.sibling, t.degree, t.depth, t.height, t.type))

コードの補足

先ほどの根付き木のコードとの変更点で少しつまずきそうな部分だけ補足説明していきたいと思います.
追加された要素として,siblingとdegreeとheightの三つがあります.
siblingとdegreeに関しては入力のleft,rightが-1である場合において注意しながら条件分岐させ,プログラムを作成することで簡単に出力の形を作ることができると思います.
heightに関しては再帰関数を用いています.
おおまかな流れとしては,呼び出し元に対して呼び出し先のheightを返し,それを+1したものが呼び出し元のheightになるようになっています.
これを再帰的に呼び出すことでrootまでの全節点のheightがわかるようになっています.

最後に

今回扱っている木構造はイメージはしやすいのですが,実装に関しては何度かやっていく中で慣れていく必要があると感じました.
実際に自分も作ってみたのは初めてだったので読みづらい部分があると思いますが,愚直に解いているので上からしっかり読み進めていけば理解できるようになっていると思います.

今後このような問題を解いてほしい,木構造のより簡単な実装などありましたらコメントよろしくお願いします!!