# tree node class Node(object): def __init__(self, ele): self.data = ele self.left = None self.right = None self.height = 1 class AvlTree(object): # insert method def insertNode(self, node, ele): if not node: return Node(ele) elif ele < node.val: node.left = self.insertNode(node.left, ele) else: node.right = self.insertNode(node.right, ele) # height updation node.height = 1 + max(self.getHeight(node.left), self.getHeight(node.right)) balance = self.getBalance(node) # ll rotation if balance > 1 and ele < node.left.val: return self.rightRotate(node) # rr rotation if balance < -1 and ele > node.right.val: return self.leftRotate(node) # lr rotation if balance > 1 and ele > node.left.val: node.left = self.leftRotate(node.left) return self.rightRotate(node) # rl rotation if balance < -1 and ele < node.right.val: node.right = self.rightRotate(node.right) return self.leftRotate(node) return node # deletion method def delete(self, node, ele): if not node: return node elif ele < node.val: node.left = self.delete(node.left, ele) elif ele > node.val: node.right = self.delete(node.right, ele) else: if node.left is None: temp = node.right node = None return temp elif node.right is None: temp = node.left node = None return temp temp = self.getMinValueNode(node.right) node.val = temp.val node.right = self.delete(node.right, temp.val) # one node case if node is None: return node # height updation node.height = 1 + max(self.getHeight(node.left), self.getHeight(node.right)) balance = self.getBalance(node) # ll rotation if balance > 1 and self.getBalance(node.left) >= 0: return self.rightRotate(node) # rr rotation if balance < -1 and self.getBalance(node.right) <= 0: return self.leftRotate(node) # lr rotation if balance > 1 and self.getBalance(node.left) < 0: node.left = self.leftRotate(node.left) return self.rightRotate(node) # rl rotation if balance < -1 and self.getBalance(node.right) > 0: node.right = self.rightRotate(node.right) return self.leftRotate(node) return node def leftRotate(self, n): x = n.right y = x.left # rotation x.left = n n.right = y n.height = 1 + max(self.getHeight(n.left), self.getHeight(n.right)) x.height = 1 + max(self.getHeight(x.left), self.getHeight(x.right)) return x def rightRotate(self, n): x = n.left y = x.right # rotation x.right = n n.left = y n.height = 1 + max(self.getHeight(n.left), self.getHeight(n.right)) x.height = 1 + max(self.getHeight(x.left), self.getHeight(x.right)) return x def height(self, node): if not node: return 0 return node.height def balFact(self, node): if not node: return 0 return self.getHeight(node.left) - self.getHeight(node.right) def getMinValueNode(self, node): if node is None or node.left is None: return node return self.getMinValueNode(node.left) def preOrderTraversal(self, node): if not node: return print("{0} ".format(node.val), end="") self.preOrder(node.left) self.preOrder(node.right) def main(): tree = AvlTree() root = None while True: print("Choose an operation :\n1. INSERTION\n2. DELETION\n3. PREORDER TRAVERSAL\n4. Exit Program") ch = int(input("\tMAKE YOUR CHOICE : ")) if ch == 1: ele = int(input("\nEnter the element for insertion : ")) tree.insertNode(root, ele) elif ch == 2: ele = int(input("\nEnter the element for deletion : ")) tree.delete(root, ele) elif ch == 3: print("\nPREORDER TRAVERSAL :", end=" ") tree.preOrderTraversal(root) elif ch == 4: exit(1) main()