Pythonで二分木の通り順走査（Inorder Traversal）によりN番目のノードを取得するプログラム

二分木の通り順走査（inorder traversal）を利用して「n番目」のノードを取得したい場合、ルート要素の設定、左右へのノード挿入、通り順走査の実行といった機能を持つ二分木クラスを作成するのが一般的です。クラスのインスタンスを生成すれば、これらのメソッドを自由に呼び出せるようになります。

本記事では、実際のコード例とともに、その仕組みと動作をわかりやすく解説します。

サンプルコード

class BinaryTree_struct:
    def __init__(self, key=None):
        self.key = key
        self.left = None
        self.right = None

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

    def inorder_nth(self, n):
        return self.inorder_nth_helper_fun(n, [])

    def inorder_nth_helper_fun(self, n, in_ord):
        if self.left is not None:
            temp = self.left.inorder_nth_helper_fun(n, in_ord)
            if temp is not None:
                return temp
        in_ord.append(self)
        if n == len(in_ord):
            return self
        if self.right is not None:
            temp = self.right.inorder_nth_helper_fun(n, in_ord)
            if temp is not None:
                return temp

    def insert_t0_left(self, new_node):
        self.left = new_node

    def insert_to_right(self, new_node):
        self.right = new_node

    def search_elem(self, key):
        if self.key == key:
            return self
        if self.left is not None:
            temp = self.left.search_elem(key)
            if temp is not None:
                return temp
        if self.right is not None:
            temp = self.right.search_elem(key)
            return temp
        return None

btree_instance = None

print('Menu (this assumes no duplicate keys)')
print('insert <data> at root')
print('insert <data> left of <data>')
print('insert <data> right of <data>')
print('inorder ')
print('quit')

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

    operation = do[0].strip().lower()
    if operation == 'insert':
        data = int(do[1])
        new_node = BinaryTree_struct(data)
        suboperation = do[2].strip().lower()
        if suboperation == 'at':
            btree_instance = new_node
        else:
            position = do[4].strip().lower()
            key = int(position)
            ref_node = None
            if btree_instance is not None:
                ref_node = btree_instance.search_elem(key)
            if ref_node is None:
                print('No such key.')
                continue
            if suboperation == 'left':
                ref_node.insert_t0_left(new_node)
            elif suboperation == 'right':
                ref_node.insert_to_right(new_node)

    elif operation == 'inorder':
        if btree_instance is not None:
            index = int(do[1].strip().lower())
            node = btree_instance.inorder_nth(index)
            if node is not None:
                print('nth term of inorder traversal: {}'.format(node.key))
            else:
                print('The index exceeds maximum possible index.')
        else:
            print('The tree is empty...')

    elif operation == 'quit':
        break

実行結果

Menu (this assumes no duplicate keys)
insert <data> at root
insert <data> left of <data>
insert <data> right of <data>
inorder
quit
What would you like to do? insert 5 at root
What would you like to do? insert 6 left of 5
What would you like to do? insert 8 right of 5
What would you like to do? inorder 5
The index exceeds maximum possible index.
What would you like to do? 6
6

コードの解説

• 必要な属性を備えた「BinaryTree_struct」クラスを定義します。

• コンストラクタ（__init__）では、左の子ノードと右の子ノードを初期値「None」に設定します。

• 「set_root」メソッドは、二分木のルート要素を設定するために使用します。

• 「inorder_nth」メソッドは、再帰処理によって通り順走査を行い、指定された番号のノードを返します。

• このため、補助的な役割を担うヘルパー関数「inorder_nth_helper_fun」が併せて定義されています。

• 「insert_to_right」メソッドは、対象ノードの右側に新しい要素を追加します。

• 「insert_t0_left」メソッドは、対象ノードの左側に新しい要素を追加します。

• 「search_elem」メソッドは、特定のキーを持つノードを木の中から検索します。

• 「BinaryTree_struct」クラスのオブジェクト（インスタンス）を生成します。

• ユーザーから実行したい操作の入力を受け付けます。

• ユーザーの選択内容に応じて、対応する操作が実行されます。

• 最後に、処理結果がコンソールに出力されます。

ポイントまとめ

このプログラムの核となるのは、通り順走査（左の子 → 自身 → 右の子 の順で訪問）を行いながら、訪問したノードをリストに記録し、そのリストの長さが指定した番号 n と一致した時点で該当ノードを返す仕組みです。また、指定したインデックスが木のノード数を超えている場合は、適切なエラーメッセージが表示されるようになっています。