''' 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)