Pythonで二分木の全ノードの合計を求めるプログラムの実装方法

木構造のすべてのノードの合計を求めたい場合、まずクラスを作成し、その中にルートノードを設定するメソッド、ツリーへ要素を追加するメソッド、特定の要素を検索するメソッド、さらにツリー内の全要素を合計して合計値を返すメソッドなどを定義します。クラスのインスタンスを生成すれば、これらのメソッドにアクセスして自由に利用できます。

本記事では、対話形式のメニューを使ってノードの追加や合計の計算が行えるサンプルプログラムを紹介します。ポイントは sum_node メソッドで、再帰呼び出しを利用して「自分自身のキーの値＋すべての子ノードの合計」を順次加算していくことで、木全体の合計をシンプルに求められる点です。計算量はノード数 n に対して O(n) となります。

サンプルコード

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

   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 sum_node(self):
      my_summation = self.key
      for child in self.children:
         my_summation = my_summation + child.sum_node()
      return my_summation

my_instance = None

print('Menu (assume no duplicate keys)')
print('add <data> at root')
print('add <data> below <data>')
print('sum')
print('quit')

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

   operation = my_input[0].strip().lower()
   if operation == 'add':
      data = int(my_input[1])
      new_node = Tree_struct(data)
      suboperation = my_input[2].strip().lower()
      if suboperation == 'at':
         my_instance = new_node
      elif suboperation == 'below':
         position = my_input[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
         ref_node.add_node(new_node)

   elif operation == 'sum':
      if my_instance is None:
         print('The tree is empty')
      else:
         my_summation = my_instance.sum_node()
         print('Sum of all nodes is: {}'.format(my_summation))

   elif operation == 'quit':
      break

実行結果

Menu (assume no duplicate keys)
add <data> at root
add <data> below <data>
sum
quit
What operation would you do ? add 5 at root
What operation would you do ? add 7 below 5
What operation would you do ? add 0 below 7
What operation would you do ? sum
Sum of all nodes is: 12
What operation would you do ? quit

コードの解説

• 必要な属性を持つ「Tree_struct」クラスを定義します。

• 「__init__」関数では、ノードのキーと、子ノードを格納するための空のリストを初期化します。

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

• 「add_node」メソッドは、ツリーに新しいノード（要素）を追加します。

• 「search_node」メソッドは、指定したキーを持つノードを再帰的に検索し、見つかったノードを返します。

• 「sum_node」メソッドは、自身のキーにすべての子ノードの合計を加算することで、木全体の合計値を求めます。

• インスタンス変数 my_instance を作成し、初期値として None を代入します。

• while ループの中で、実行したい操作をユーザー入力として受け取ります。

• 入力されたコマンド（add / sum / quit）に応じて、対応する処理が実行されます。

• 処理結果はコンソールに出力され、「quit」が入力されるとプログラムを終了します。