Python Program to Count Number of Non Leaf Nodes of a given Tree

When it is required to find the count of the non leaf nodes in a Tree, a ‘Tree_structure’ class is created, methods to set a root value, and to add other values are defined. 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_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

Output

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

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 ‘count_non_leaf_nodes’ is defined, that helps get the count of the non leaf nodes of the Tree.

• It is a recursive function.

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

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

• This output is displayed on the console.