Pythonで木構造の非リーフノード（葉以外のノード）の数をカウントする方法

木構造（ツリー）において、リーフノード（葉ノード）以外、つまり子ノードを持つノードの数を求めたい場合があります。そのようなときは、「Tree_structure」というクラスを作成し、ルートの値を設定するメソッドや、新しい値を追加するメソッドなどを定義します。さらに、ユーザーが選択できる複数の操作メニューを用意し、選択内容に応じて木構造に対する処理を実行できるようにします。

以下に具体的な実装例を示します。

サンプルコード

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

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

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

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

    def count_non_leaf_node(self):
        nonleaf_count = 0
        if self.children != []:
            nonleaf_count = 1
        for child in self.children:
            nonleaf_count = nonleaf_count + child.count_non_leaf_node()
        return nonleaf_count

tree = None

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

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

    operation = my_input[0].strip().lower()
    if operation == 'add':
        data = int(my_input[1])
        newNode = Tree_structure(data)
        suboperation = my_input[2].strip().lower()
        if suboperation == 'at':
            tree = newNode
        elif suboperation == 'below':
            position = my_input[3].strip().lower()
            key = int(position)
            ref_node = None
            if tree is not None:
                ref_node = tree.search_val(key)
            if ref_node is None:
                print('No such key.')
                continue
            ref_node.add_vals(newNode)

    elif operation == 'count':
        if tree is None:
            print('The tree is empty ')
        else:
            count = tree.count_non_leaf_node()
            print('The number of non-leaf nodes are : {}'.format(count))

    elif operation == 'quit':
        break

実行結果

Menu (this assumes no duplicate keys)
add <data> at root
add <data> below <data>
count
quit
What operation would you like to perform ? add 34 at root
What operation would you like to perform ? add 78 below 34
What operation would you like to perform ? add 56 below 78
What operation would you like to perform ? add 90 below 56
What operation would you like to perform ? count
The number of non-leaf nodes are : 3
What operation would you like to perform ? quit

コードの解説

• まず「Tree_structure」クラスを作成します。

• コンストラクタでは「key」にデータを格納し、「children」には空のリストを設定して、木の子ノードを管理できるようにします。

• 「set_root」メソッドは、木のルート（根）の値を設定するために使用されます。

• 「add_vals」メソッドは、木に新しい要素（ノード）を追加するために定義されています。

• 「search_val」メソッドは、指定したキーの値を持つノードを木の中から探索するためのメソッドです。

• 「count_non_leaf_node」メソッドは、木に含まれる非リーフノード（子を持つノード）の数を取得するためのメソッドです。

• このメソッドは再帰的に呼び出されることで、すべての子ノードをたどりながら合計数を計算します。

• メニューとして「ルートに追加（add at root）」「指定ノードの下に追加（add below）」「カウント（count）」「終了（quit）」の4つの選択肢を用意しています。

• ユーザーが入力した選択肢に応じて、対応する処理が実行されます。

• 実行結果はコンソールに出力されます。