Python Program to Find the Sum of all Nodes in a Tree

When it is required to get the sum of all the nodes in a Tree, a ‘Tree_structure’ class is created, methods to set a root value, and to add other values are defined. It also has a method to determine the sum of all elements of a Tree structure. Various options are given that the user can select. Based on the user’s choice, the operation is performed on the Tree elements.

Below is a demonstration of the same −

Example

 Live Demo

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

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

   def add_values(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(key)
         if temp is not None:
            return temp
      return None

   def summation_nodes(self):
      sum_val = self.key
      for child in self.children:
         sum_val = sum_val + child.summation_nodes()
      return sum_val

tree = None

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

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

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

   elif operation == 'summation':
      if tree is None:
         print('The tree is empty')
      else:
         summation_val = tree.summation_nodes()
         print('Sum of all the nodes is : {}'.format(summation_val))

   elif operation == 'quit':
      break

Output

Menu (no duplicate keys allowed)
add <data> at root
add <data> below <data>
summation
quit
What would you like to do? add 56 at root
What would you like to do? add 45 below 56
What would you like to do? add 23 below 56
What would you like to do? summation
Sum of all the nodes is : 124
What would you like to do?

Explanation

• The ‘Tree_structure’ class is created.

• It sets the ‘key’ to True and sets an empty list to children of tree.

• It has a ‘set_root’ function that helps set the root value for the Tree.

• A method named ‘add_vals’ is defined, that helps add an element to the Tree.

• Another method named ‘search_val’ is defined, that helps search for an element in the Tree.

• Another method named ‘summation_nodes’ is defined, that helps get the sum of all elements/nodes of the Tree.

• It is a recursive function.

• Four options are given, such as ‘add at root’, ‘add below’, ‘Summation’ and ‘quit’.

• Depending on the option given by user, the respective operation is performed.

• This output is displayed on the console.