Pythonでツリー（木構造）を構築し、挿入・削除・表示を操作する方法

ツリー（木構造）を構築し、要素の挿入・削除・表示といった操作を行いたい場合、必要なメソッドを備えたクラスを定義するのが一般的です。クラスのインスタンスを生成し、それを通じてノードへのアクセスや各種操作を実行します。

この記事では、対話型メニューを備えたPythonのサンプルプログラムを紹介し、その仕組みをわかりやすく解説します。

サンプルコード

class Tree_struct:
    def __init__(self, data=None, parent=None):
        self.key = data
        self.children = []
        self.parent = parent

    def set_root(self, data):
        self.key = data

    def add_node(self, node):
        self.children.append(node)

    def search_node(self, key):
        if self.key == key:
            return self
        for child in self.children:
            temp = child.search_node(key)
            if temp is not None:
                return temp
        return None

    def remove_node(self):
        parent = self.parent
        index = parent.children.index(self)
        parent.children.remove(self)
        for child in reversed(self.children):
            parent.children.insert(index, child)
            child.parent = parent

    def bfs(self):
        queue = [self]
        while queue != []:
            popped = queue.pop(0)
            for child in popped.children:
                queue.append(child)
            print(popped.key, end=' ')

my_instance = None

print('Menu (this assumes no duplicate keys)')
print('add <data> at root')
print('add <data> below <data>')
print('remove <data>')
print('display')
print('quit')

while True:
    do = input('What would you like to do? ').split()

    operation = do[0].strip().lower()
    if operation == 'add':
        data = int(do[1])
        new_node = Tree_struct(data)
        suboperation = do[2].strip().lower()
        if suboperation == 'at':
            my_instance = new_node
        elif suboperation == 'below':
            position = do[3].strip().lower()
            key = int(position)
            ref_node = None
            if my_instance is not None:
                ref_node = my_instance.search_node(key)
            if ref_node is None:
                print('No such key.')
                continue
            new_node.parent = ref_node
            ref_node.add_node(new_node)

    elif operation == 'remove':
        data = int(do[1])
        to_remove = my_instance.search_node(data)
        if my_instance == to_remove:
            if my_instance.children == []:
                my_instance = None
            else:
                leaf = my_instance.children[0]
                while leaf.children != []:
                    leaf = leaf.children[0]
                leaf.parent.children.remove_node(leaf)
                leaf.parent = None
                leaf.children = my_instance.children
                my_instance = leaf
        else:
            to_remove.remove_node()

    elif operation == 'display':
        if my_instance is not None:
            print('Breadth First Search traversal is : ', end='')
            my_instance.bfs()
            print()
        else:
            print('The tree is empty')

    elif operation == 'quit':
        break

実行結果

Menu (this assumes no duplicate keys)
add <data> at root
add <data> below <data>
remove <data>
display
quit
What would you like to do? add 5 at root
What would you like to do? add 6 below 5
What would you like to do? add 8 below 6
What would you like to do? remove 8
What would you like to do? display
Breadth First Search traversal is : 5 6
What would you like to do? quit

コードの解説

• 必要な属性を持つ「Tree_struct」クラスを定義しています。
• 「__init__」コンストラクタは、ノードのキー・子ノードのリスト・親ノードを初期化します。
• 「set_root」メソッドは、ツリーのルートとなる値を設定するために用意されています。
• 「add_node」メソッドは、ツリーに新しいノードを追加します。
• 「search_node」メソッドは、再帰的に処理を行い、指定されたキーを持つノードを検索します。
• 「remove_node」メソッドは、ツリーからノードを削除します。削除対象ノードの子は、その親ノードへ引き継がれます。
• 「bfs」メソッドは、キューを利用してツリーに対して幅優先探索（BFS）を実行し、各ノードのキーを順番に出力します。
• インスタンス変数「my_instance」が作成され、初期状態では「None」が代入されています。
• ユーザーから実行したい操作の入力を受け付けます。
• ユーザーの選択に応じて、対応する操作が実行されます。
• 処理結果がコンソールに表示されます。

メニューで使えるコマンド一覧

• add <データ> at root … 入力した値をルートとして新しいツリーを作成します。
• add <データ> below <親の値> … 指定した親ノードの下に新しい子ノードを追加します。該当するキーが存在しない場合は「No such key.」と表示されます。
• remove <データ> … 指定したキーのノードを削除します。
• display … 幅優先探索（BFS）による走査結果を表示します。ツリーが空の場合は「The tree is empty」と出力されます。
• quit … プログラムを終了します。

なお、メニューの冒頭にも表示されている通り、このプログラムはキーの重複が存在しないことを前提として設計されています。同じキーを複数登録すると、検索や削除の動作が意図どおりにならない場合があるため注意してください。