Python Program to Construct a Tree & Perform Insertion, Deletion, Display

When it is required to construct a binary tree, and perform operations such as inserting an element, deleting an element and displaying elements of the tree, a class is defined with methods in it. An instance of the class is defined and this is used to access the elements and perform operations.

Below is a demonstration of the same −

Example

 Live Demo

class Tree_struct:    def __init__(self, data=None, parent=None):       self.key = data       self.children = []       self.parent = parent    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 remove_node(self):       parent = self.parent       index = parent.children.index(self)       parent.children.remove(self)       for child in reversed(self.children):          parent.children.insert(index, child)          child.parent = parent    def bfs(self):       queue = [self]       while queue != []:          popped = queue.pop(0)          for child in popped.children:             queue.append(child)          print(popped.key, end=' ') my_instance = None print('Menu (this assumes no duplicate keys)') print('add <data> at root') print('add <data> below <data>') print('remove <data>') print('display') print('quit') while True:    do = input('What would you like to do? ').split()    operation = do[0].strip().lower()    if operation == 'add':       data = int(do[1])       new_node = Tree_struct(data)       suboperation = do[2].strip().lower()       if suboperation == 'at':          my_instance = new_node       elif suboperation == 'below':          position = do[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          new_node.parent = ref_node          ref_node.add_node(new_node)    elif operation == 'remove':       data = int(do[1])       to_remove = my_instance.search_node(data)       if my_instance == to_remove:          if my_instance.children == []:             my_instance = None          else:             leaf = my_instance.children[0]             while leaf.children != []:                leaf = leaf.children[0]             leaf.parent.children.remove_node(leaf)             leaf.parent = None             leaf.children = my_instance.children             my_instance = leaf       else:          to_remove.remove_node()    elif operation == 'display':       if my_instance is not None:          print('Breadth First Search traversal is : ', end='')          my_instance.bfs()          print()       else:          print('The tree is empty')    elif operation == 'quit':       break

Output

Menu (this assumes no duplicate keys) add <data> at root add <data> below <data> remove <data> display quit What would you like to do? add 5 at root What would you like to do? add 6 below 5 What would you like to do? add 8 below 6 What would you like to do? remove 8 What would you like to do? display Breadth First Search traversal is : 5 6 What would you like to 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.

• Another method named ‘set_root’ is defined to specify the root of the tree.

• Another method named ‘add_node’ is defined that helps add nodes to the tree.

• Another method named ‘search_node’ is defined that helps search for a specific element.

• A method named ‘remove_node’ is defined, that deletes elements from the tree.

• Another method named ‘bfs’ is defined, that helps perform breadth first search on the tree.

• 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.