Merge pull request keon#244 from muneeb20/add_avl_tree_implementation

Avl tree implementation (keon#198)

tree/avl/avl.py

@@ -0,0 +1,124 @@
+from tree.tree import TreeNode
+
+
+class AvlTree(object):
+    """
+    An avl tree.
+    """
+
+    def __init__(self):
+        # Root node of the tree.
+        self.node = None
+        self.height = -1
+        self.balance = 0
+
+    def insert(self, key):
+        """
+        Insert new key into node
+        """
+        # Create new node
+        n = TreeNode(key)
+        if not self.node:
+            self.node = n
+            self.node.left = AvlTree()
+            self.node.right = AvlTree()
+        elif key < self.node.val:
+            self.node.left.insert(key)
+        elif key > self.node.val:
+            self.node.right.insert(key)
+        self.re_balance()
+
+    def re_balance(self):
+        """
+        Re balance tree. After inserting or deleting a node,
+        """
+        self.update_heights(recursive=False)
+        self.update_balances(False)
+
+        while self.balance < -1 or self.balance > 1:
+            if self.balance > 1:
+                if self.node.left.balance < 0:
+                    self.node.left.rotate_left()
+                    self.update_heights()
+                    self.update_balances()
+                self.rotate_right()
+                self.update_heights()
+                self.update_balances()
+
+            if self.balance < -1:
+                if self.node.right.balance > 0:
+                    self.node.right.rotate_right()
+                    self.update_heights()
+                    self.update_balances()
+                self.rotate_left()
+                self.update_heights()
+                self.update_balances()
+
+    def update_heights(self, recursive=True):
+        """
+        Update tree height
+        """
+        if self.node:
+            if recursive:
+                if self.node.left:
+                    self.node.left.update_heights()
+                if self.node.right:
+                    self.node.right.update_heights()
+
+            self.height = 1 + max(self.node.left.height, self.node.right.height)
+        else:
+            self.height = -1
+
+    def update_balances(self, recursive=True):
+        """
+        Calculate tree balance factor
+
+        """
+        if self.node:
+            if recursive:
+                if self.node.left:
+                    self.node.left.update_balances()
+                if self.node.right:
+                    self.node.right.update_balances()
+
+            self.balance = self.node.left.height - self.node.right.height
+        else:
+            self.balance = 0
+
+    def rotate_right(self):
+        """
+        Right rotation
+        """
+        new_root = self.node.left.node
+        new_left_sub = new_root.right.node
+        old_root = self.node
+
+        self.node = new_root
+        old_root.left.node = new_left_sub
+        new_root.right.node = old_root
+
+    def rotate_left(self):
+        """
+        Left rotation
+        """
+        new_root = self.node.right.node
+        new_left_sub = new_root.left.node
+        old_root = self.node
+
+        self.node = new_root
+        old_root.right.node = new_left_sub
+        new_root.left.node = old_root
+
+    def in_order_traverse(self):
+        """
+        In-order traversal of the tree
+        """
+        result = []
+
+        if not self.node:
+            return result
+
+        result.extend(self.node.left.in_order_traverse())
+        result.append(self.node.key)
+        result.extend(self.node.right.in_order_traverse())
+        return result