Find the Sum of All Nodes in a Binary Tree using Python

When it is required to find the sum of all the nodes of a tree, a class is created, and it contains methods to set the root node, add elements to the tree, search for a specific element, and add elements of the tree to find the sum and so on. An instance of the class can be created to access and use these methods.

Below is a demonstration of the same −

Example

 Live Demo

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

Output

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

Explanation

• The ‘Tree_struct’ class with required attributes is created.

• It has an ‘init’ function that is used to create an empty list.

• A ‘set_root’ method is defined that helps set the root value of the binary tree.

• It has an ‘add_node’ method that helps add elements to the tree.

• A method named ‘search_elem’ is defined, that helps search for a specific element.

• A method named ‘sum_node is defined, that helps add elements of the tree and find the sum.

• An instance is created and assigned to ‘None’.

• The user input is taken for the operation that needs to be performed.

• Depending on the user’ choice, the operation is performed.

• Relevant output is displayed on the console.