AVL Tree

code/AVL.py

@@ -0,0 +1,180 @@
+'''
+A AVL tree
+'''
+
+
+class Node:
+
+    def __init__(self, label):
+        self.label = label
+        self._parent = None
+        self._left = None
+        self._right = None
+        self.height = 0
+
+    @property
+    def right(self):
+        return self._right
+
+    @right.setter
+    def right(self, node):
+        if node is not None:
+            node._parent = self
+            self._right = node
+
+    @property
+    def left(self):
+        return self._left
+
+    @left.setter
+    def left(self, node):
+        if node is not None:
+            node._parent = self
+            self._left = node
+
+    @property
+    def parent(self):
+        return self._parent
+
+    @parent.setter
+    def parent(self, node):
+        if node is not None:
+            self._parent = node
+            self.height = self.parent.height + 1
+        else:
+            self.height = 0
+
+
+class AVL:
+
+    def __init__(self):
+        self.root = None
+        self.size = 0
+
+    def insert(self, value):
+        node = Node(value)
+
+        if self.root is None:
+            self.root = node
+            self.root.height = 0
+            self.size = 1
+        else:
+            # Same as Binary Tree
+            dad_node = None
+            curr_node = self.root
+
+            while True:
+                if curr_node is not None:
+
+                    dad_node = curr_node
+
+                    if node.label < curr_node.label:
+                        curr_node = curr_node.left
+                    else:
+                        curr_node = curr_node.right
+                else:
+                    node.height = dad_node.height
+                    dad_node.height += 1
+                    if node.label < dad_node.label:
+                        dad_node.left = node
+                    else:
+                        dad_node.right = node
+                    self.rebalance(node)
+                    self.size += 1
+                    break
+
+    def rebalance(self, node):
+        n = node
+
+        while n is not None:
+            height_right = n.height
+            height_left = n.height
+
+            if n.right is not None:
+                height_right = n.right.height
+
+            if n.left is not None:
+                height_left = n.left.height
+
+            if abs(height_left - height_right) > 1:
+                if height_left > height_right:
+                    left_child = n.left
+                    if left_child is not None:
+                        h_right = (right_child.right.height
+                                    if (right_child.right is not None) else 0)
+                        h_left = (right_child.left.height
+                                    if (right_child.left is not None) else 0)
+                    if (h_left > h_right):
+                        self.rotate_left(n)
+                        break
+                    else:
+                        self.double_rotate_right(n)
+                        break
+                else:
+                    right_child = n.right
+                    if right_child is not None:
+                        h_right = (right_child.right.height
+                            if (right_child.right is not None) else 0)
+                        h_left = (right_child.left.height
+                            if (right_child.left is not None) else 0)
+                    if (h_left > h_right):
+                        self.double_rotate_left(n)
+                        break
+                    else:
+                        self.rotate_right(n)
+                        break
+            n = n.parent
+
+    def rotate_left(self, node):
+        aux = node.parent.label
+        node.parent.label = node.label
+        node.parent.right = Node(aux)
+        node.parent.right.height = node.parent.height + 1
+        node.parent.left = node.right
+
+
+    def rotate_right(self, node):
+        aux = node.parent.label
+        node.parent.label = node.label
+        node.parent.left = Node(aux)
+        node.parent.left.height = node.parent.height + 1
+        node.parent.right = node.right
+
+    def double_rotate_left(self, node):
+        self.rotate_right(node.getRight().getRight())
+        self.rotate_left(node)
+
+    def double_rotate_right(self, node):
+        self.rotate_left(node.getLeft().getLeft())
+        self.rotate_right(node)
+
+    def empty(self):
+        if self.root is None:
+            return True
+        return False
+
+    def preShow(self, curr_node):
+        if curr_node is not None:
+            self.preShow(curr_node.left)
+            print(curr_node.label, end=" ")
+            self.preShow(curr_node.right)
+
+    def preorder(self, curr_node):
+        if curr_node is not None:
+            self.preShow(curr_node.left)
+            self.preShow(curr_node.right)
+            print(curr_node.label, end=" ")
+
+    def getRoot(self):
+        return self.root
+
+t = AVL()
+t.insert(1)
+t.insert(2)
+t.insert(3)
+# t.preShow(t.root)
+# print("\n")
+# t.insert(4)
+# t.insert(5)
+# t.preShow(t.root)
+# t.preorden(t.root)