Python-Practicle/BST_del_node.py at master · pavanpatil1/Python-Practicle · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
from random import randint

class Node(object):
    def __init__(self, key):
        self.key   = key
        self.left  = None
        self.right = None

class BST(object):
    def __init__(self):
        self.root = None

    def insert(self, key, parent=None):
        if parent is None:
            parent = self.root
        if self.root is None:
            self.root = Node(key)
            print("Added", key, "as root node")
        else:
            if key > parent.key:                                            # Traverse the right sub-tree if key > current node
                if parent.right is None:
                    print("\nAdding", key,"to", parent.key, "\'s right")
                    parent.right = Node(key)
                else:
                    parent = parent.right
                    self.insert(key, parent)
            else:                                                           # Traverse the left sub-tree if key < current node
                if parent.left is None:
                    print("\nAdding", key,"to", parent.key, "\'s left")
                    parent.left = Node(key)
                else:
                    parent = parent.left
                    self.insert(key, parent)

    def remove(self, key, parent=None):
        if self.root is None:
            return None
        elif parent is None:
            parent = self.root

        if key < parent.key:
            if parent.left is None:
                print("Not present")
            else:
                parent.left = self.remove(key, parent.left)             # if key < current node, traverse the left sub-tree
        elif key > parent.key:
            if parent.right is None:
                print("Not present")
            else:
                parent.right = self.remove(key, parent.right)           # if key > current node, traverse the right sub-tree
        elif key == parent.key:                                                       # key found
            if parent.left is None:                                 # base case: if node has no child or 1 child
                temp = parent.right
                parent = None
                return temp

            elif parent.right is None:
                temp = parent.left
                parent = None
                return temp

            else:                                                   # if node has 2 child, replace key to be removed with its successor
                temp = self.successor(parent)
                parent.key = temp.key
                parent.right = self.remove(temp.key, parent.right)
        return parent

    def search(self, key, parent=None):
        if parent is None:
            if self.root is None:
                return None
            if self.root == key:
                return self.root
            parent = self.root
        if parent.key == key:
            return parent
        if key > parent.key:
            if parent.right is None:
                return None
            return self.search(key, parent.right)
        elif key < parent.key:
            if parent.left is None:
                return None
            return self.search(key, parent.left)

    def successor(self, key):
        '''method to find the successor of the given key'''
        if isinstance(key, Node):
            node = self.search(key.key)
        else:
            node = self.search(key)
        succ = node.right
        if succ is not None:
            while succ.left:
                succ = succ.left
        return succ

    def inorder(self, node=1):
        '''method to print BST inorder (left, root, right)'''
        if node == 1:
            node = self.root
            print("\nBST is: (inorder)\n")
        if node:
            self.inorder(node.left)
            print(node.key, end=", ")
            self.inorder(node.right)


def main():
    bst = BST()
    height = 4
    n = 2**height - 1
    for _ in range(n):
        bst.insert(randint(1, 100))
    bst.inorder()

    bst.remove(int(input("\n")))
    bst.inorder()
    bst.remove(int(input("\n")))
    bst.inorder()


if __name__ == "__main__":
    main()