Python Program to Construct a Tree and Perform Tree Operations

«
»

This is a Python program to construct a tree and perform insertion, deletion and display operations.

Problem Description

The program creates a tree and presents a menu to the user to perform insertion, deletion and display operations.

Problem Solution

1. Create a class Tree with instance variables key and children.
2. Define methods set_root, add, remove, bfs_display and search.
3. The method set_root takes a key as argument and sets the instance variable key equal to it.
4. The method add appends a node to the list children.
5. The method search returns a node with a specified key.
6. The method remove removes the current node from the tree.
7. The method bfs_display displays the BFS traversal of the tree.

Program/Source Code

Here is the source code of a Python program to construct a tree and perform insertion, deletion and display operations. The program output is shown below.

class Tree:
    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(self, node):
        self.children.append(node)
 
    def search(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 remove(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_display(self):
        queue = [self]
        while queue != []:
            popped = queue.pop(0)
            for child in popped.children:
                queue.append(child)
            print(popped.key, end=' ')
 
 
tree = 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(data)
        suboperation = do[2].strip().lower() 
        if suboperation == 'at':
            tree = new_node
        elif suboperation == 'below':
            position = do[3].strip().lower()
            key = int(position)
            ref_node = None
            if tree is not None:
                ref_node = tree.search(key)
            if ref_node is None:
                print('No such key.')
                continue
            new_node.parent = ref_node
            ref_node.add(new_node)
 
    elif operation == 'remove':
        data = int(do[1])
        to_remove = tree.search(data)
        if tree == to_remove:
            if tree.children == []:
                tree = None
            else:
                leaf = tree.children[0]
                while leaf.children != []:
                    leaf = leaf.children[0]
                leaf.parent.children.remove(leaf)
                leaf.parent = None
                leaf.children = tree.children
                tree = leaf
        else:
            to_remove.remove()
 
    elif operation == 'display':
        if tree is not None:
            print('BFS traversal display: ', end='')
            tree.bfs_display()
            print()
        else:
            print('Tree is empty.')
 
    elif operation == 'quit':
        break
Program Explanation

1. A variable is created to store the binary tree.
2. The user is presented with a menu to perform operations on the tree.
3. The corresponding methods are called to perform each operation.
4. If the node to be removed is the root node, the removal is carried out by finding a leaf node and replacing the root with it. If the node to be removed is not the root node then the remove method is called on the node.

Note: Join free Sanfoundry classes at Telegram or Youtube
advertisement
advertisement
Runtime Test Cases
Case 1:
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 1 at root
What would you like to do? add 2 below 1
What would you like to do? add 3 below 1
What would you like to do? add 4 below 2
What would you like to do? add 5 below 2
What would you like to do? display
BFS traversal display: 1 2 3 4 5 
What would you like to do? remove 1
What would you like to do? display
BFS traversal display: 4 2 3 5 
What would you like to do? remove 5
What would you like to do? display
BFS traversal display: 4 2 3 
What would you like to do? remove 4
What would you like to do? display
BFS traversal display: 2 3 
What would you like to do? remove 3
What would you like to do? remove 2
What would you like to do? display
Tree is empty.
What would you like to do? quit
 
Case 2:
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 7 below 5
What would you like to do? add 9 below 7
What would you like to do? add 11 below 9
What would you like to do? add 12 below 7
What would you like to do? display
BFS traversal display: 5 7 9 12 11 
What would you like to do? remove 9
What would you like to do? display
BFS traversal display: 5 7 11 12 
What would you like to do? remove 12
What would you like to do? display
BFS traversal display: 5 7 11 
What would you like to do? remove 7
What would you like to do? display
BFS traversal display: 5 11 
What would you like to do? remove 5
What would you like to do? display
BFS traversal display: 11 
What would you like to do? quit

Sanfoundry Global Education & Learning Series – Python Programs.

To practice all Python programs, here is complete set of 150+ Python Problems and Solutions.

Take Python Programming Mock Tests - Chapterwise!
Start the Test Now: Chapter 1, 2, 3, 4, 5, 6, 7, 8, 9, 10

advertisement
advertisement
Subscribe to our Newsletters (Subject-wise). Participate in the Sanfoundry Certification contest to get free Certificate of Merit. Join our social networks below and stay updated with latest contests, videos, internships and jobs!

Youtube | Telegram | LinkedIn | Instagram | Facebook | Twitter | Pinterest
Manish Bhojasia - Founder & CTO at Sanfoundry
Manish Bhojasia, a technology veteran with 20+ years @ Cisco & Wipro, is Founder and CTO at Sanfoundry. He lives in Bangalore, and focuses on development of Linux Kernel, SAN Technologies, Advanced C, Data Structures & Alogrithms. Stay connected with him at LinkedIn.

Subscribe to his free Masterclasses at Youtube & technical discussions at Telegram SanfoundryClasses.