Create avltree.py

Author: veerendra-poosala

avltree.py

@@ -0,0 +1,195 @@
+from __future__ import print_function
+#creating avl tree node
+
+class treenode(object):
+    def __init__(self,key):
+        self.key=key
+        self.left=None
+        self.right=None
+        self.height=1
+
+#creating avl tree class
+class avltree(object):
+    def get_height(self,root):
+        if not root:
+            return 0
+        return root.height
+
+    def get_balance(self,root):
+        if not root:
+            return 0
+        return self.get_height(root.left)-self.get_height(root.right)
+
+    def insert(self,root,key):
+        #perform normal bst
+        if not root:
+            return treenode(key)
+        elif key<root.key:
+            root.left=self.insert(root.left,key)
+        else:
+            root.right=self.insert(root.right,key)
+        #step 2- update the height of the ancestor node
+        root.height=1+max(self.get_height(root.left),self.get_height(root.right))
+        #3.get the balance factor
+        balance=self.get_balance(root)
+        #4.if the node is unbalanced ,then try out 4 cases
+        #LL
+        if balance >1 and key < root.left.key:
+            return self.righrrotate(root)
+        #RR
+        if balance < -1 and key>root.right.key:
+            return self.leftrotate(root)
+        #LR
+        if balance >1 and key > root.left.key:
+            return self.righrrotate(root)
+        #RL
+        if balance <-1 and key<root.right.key:
+            return self.leftrotate(root)
+        return root
+
+    # Recursive function to delete a node with
+    # given key from subtree with given root.
+    # It returns root of the modified subtree.
+    def delete(self, root, key):
+  
+        # Step 1 - Perform standard BST delete
+        if not root:
+            return root
+  
+        elif key < root.key:
+            root.left = self.delete(root.left, key)
+  
+        elif key > root.key:
+            root.right = self.delete(root.right, key)
+  
+        else:
+            if root.left is None:
+                temp = root.right
+                root = None
+                return temp
+  
+            elif root.right is None:
+                temp = root.left
+                root = None
+                return temp
+  
+            temp = self.getMinValueNode(root.right)
+            root.key = temp.key
+            root.right = self.delete(root.right,
+                                      temp.key)
+  
+        # If the tree has only one node,
+        # simply return it
+        if root is None:
+            return root
+  
+        # Step 2 - Update the height of the 
+        # ancestor node
+        root.height = 1 + max(self.get_height(root.left),
+                            self.get_height(root.right))
+  
+        # Step 3 - Get the balance factor
+        balance = self.get_balance(root)
+  
+        # Step 4 - If the node is unbalanced, 
+        # then try out the 4 cases
+        # Case 1 - Left Left
+        if balance > 1 and self.get_balance(root.left) >= 0:
+            return self.righrrotate(root)
+  
+        # Case 2 - Right Right
+        if balance < -1 and self.get_balance(root.right) <= 0:
+            return self.leftrotate(root)
+  
+        # Case 3 - Left Right
+        if balance > 1 and self.get_balance(root.left) < 0:
+            root.left = self.leftRotate(root.left)
+            return self.righrrotate(root)
+  
+        # Case 4 - Right Left
+        if balance < -1 and self.get_balance(root.right) > 0:
+            root.right = self.rightRotate(root.right)
+            return self.leftrotate(root)
+  
+        return root
+    
+    def leftrotate(self,z):
+        y=z.right
+        t2=y.left
+        #perform rotation
+        y.left=z
+        z.right=t2
+        #upadate heights
+        z.height=1+max(self.get_height(z.left),self.get_height(z.right))
+        y.height=1+max(self.get_height(y.left),self.get_height(y.right))
+        return y
+
+    def righrrotate(self,z):
+        y=z.left
+        t2=y.right
+        #perform rotation
+        y.right=z
+        z.left=t2
+        #upadate heights
+        z.height=1+max(self.get_height(z.left),self.get_height(z.right))
+        y.height=1+max(self.get_height(y.left),self.get_height(y.right))
+        return y
+    def preOrder(self, root):
+ 
+        if not root:
+            return
+ 
+        print("{0} ".format(root.key), end="")
+        self.preOrder(root.left)
+        self.preOrder(root.right)
+    def getMinValueNode(self, root):
+        if root is None or root.left is None:
+            return root
+  
+        return self.getMinValueNode(root.left)
+  
+ 
+
+
+root=None
+myroot=avltree()
+root=myroot.insert(root,8)
+root=myroot.insert(root,10)
+root=myroot.insert(root,1)
+root=myroot.insert(root,6)
+root=myroot.insert(root,0)
+root=myroot.insert(root,11)
+# Preorder Traversal
+print("Preorder traversal of the",
+      "constructed AVL tree is")
+myroot.preOrder(root)
+print()
+root=myroot.insert(root,9)
+print("Preorder traversal of the",
+      "constructed AVL tree is")
+myroot.preOrder(root)
+print()
+root=myroot.insert(root,14)
+print("Preorder traversal of the",
+      "constructed AVL tree is")
+myroot.preOrder(root)
+print()
+
+myTree =avltree()
+root = None
+nums = [9, 5, 10, 0, 6, 11, -1, 1, 2]
+  
+for num in nums:
+    root = myTree.insert(root, num)
+print("Preorder traversal of the",
+      "constructed AVL my tree is")
+myTree.preOrder(root)
+print() 
+#Delete
+key = 10
+root = myTree.delete(root, key)
+  
+# Preorder Traversal
+print("Preorder Traversal after deletion -")
+myTree.preOrder(root)
+print()