""" File: avl_tree.py Created Time: 2022-12-20 Author: a16su (lpluls001@gmail.com) """ import sys from pathlib import Path sys.path.append(str(Path(__file__).parent.parent)) from modules import TreeNode, print_tree class AVLTree: """AVL tree""" def __init__(self): """Constructor""" self._root = None def get_root(self) -> TreeNode | None: """Get the root of the binary tree""" return self._root def height(self, node: TreeNode | None) -> int: """Get the node height""" # An empty subtree has height -1; a leaf has height 0 if node is not None: return node.height return -1 def update_height(self, node: TreeNode | None): """Update the node height""" # The node height equals the height of its taller subtree + 1 node.height = max([self.height(node.left), self.height(node.right)]) + 1 def balance_factor(self, node: TreeNode | None) -> int: """Get the balance factor""" # An empty subtree has balance factor 0 if node is None: return 0 # Node balance factor = left subtree height - right subtree height return self.height(node.left) - self.height(node.right) def right_rotate(self, node: TreeNode | None) -> TreeNode | None: """Right rotation""" child = node.left grand_child = child.right # Rotate node to the right around child child.right = node node.left = grand_child # Update the node height self.update_height(node) self.update_height(child) # Return the root of the rotated subtree return child def left_rotate(self, node: TreeNode | None) -> TreeNode | None: """Left rotation""" child = node.right grand_child = child.left # Rotate node to the left around child child.left = node node.right = grand_child # Update the node height self.update_height(node) self.update_height(child) # Return the root of the rotated subtree return child def rotate(self, node: TreeNode | None) -> TreeNode | None: """Perform rotations to restore balance in this subtree""" # Get the balance factor of node balance_factor = self.balance_factor(node) # Left-heavy tree if balance_factor > 1: if self.balance_factor(node.left) >= 0: # Rotate right return self.right_rotate(node) else: # Rotate left, then right node.left = self.left_rotate(node.left) return self.right_rotate(node) # Right-heavy tree elif balance_factor < -1: if self.balance_factor(node.right) <= 0: # Rotate left return self.left_rotate(node) else: # Rotate right, then left node.right = self.right_rotate(node.right) return self.left_rotate(node) # The tree is balanced; return without rotating return node def insert(self, val): """Insert a node""" self._root = self.insert_helper(self._root, val) def insert_helper(self, node: TreeNode | None, val: int) -> TreeNode: """Insert a node recursively (helper method)""" if node is None: return TreeNode(val) # 1. Find the insertion position and insert the node if val < node.val: node.left = self.insert_helper(node.left, val) elif val > node.val: node.right = self.insert_helper(node.right, val) else: # Return without inserting a duplicate node return node # Update the node height self.update_height(node) # 2. Perform rotations to restore balance in this subtree return self.rotate(node) def remove(self, val: int): """Remove a node""" self._root = self.remove_helper(self._root, val) def remove_helper(self, node: TreeNode | None, val: int) -> TreeNode | None: """Remove a node recursively (helper method)""" if node is None: return None # 1. Find and remove the node if val < node.val: node.left = self.remove_helper(node.left, val) elif val > node.val: node.right = self.remove_helper(node.right, val) else: if node.left is None or node.right is None: child = node.left or node.right # Number of children = 0: remove node and return if child is None: return None # Number of children = 1: remove node directly else: node = child else: # Number of children = 2: remove the inorder successor and use its value to replace the current node value temp = node.right while temp.left is not None: temp = temp.left node.right = self.remove_helper(node.right, temp.val) node.val = temp.val # Update the node height self.update_height(node) # 2. Perform rotations to restore balance in this subtree return self.rotate(node) def search(self, val: int) -> TreeNode | None: """Find a node""" cur = self._root # Search in a loop, stopping if cur becomes None while cur is not None: # If present, the target node lies in the right subtree of cur if cur.val < val: cur = cur.right # If present, the target node lies in the left subtree of cur elif cur.val > val: cur = cur.left # The target node was found; exit the loop else: break # Return the target node, or None if it was not found return cur """Driver code""" if __name__ == "__main__": def test_insert(tree: AVLTree, val: int): tree.insert(val) print("\nAVL tree after inserting node {}:".format(val)) print_tree(tree.get_root()) def test_remove(tree: AVLTree, val: int): tree.remove(val) print("\nAVL tree after removing node {}:".format(val)) print_tree(tree.get_root()) # Initialize an empty AVL tree avl_tree = AVLTree() # Insert a node # Observe how the AVL tree maintains balance after a node is inserted for val in [1, 2, 3, 4, 5, 8, 7, 9, 10, 6]: test_insert(avl_tree, val) # Insert a duplicate node test_insert(avl_tree, 7) # Remove a node # Observe how the AVL tree maintains balance after a node is removed test_remove(avl_tree, 8) # Remove a node with 0 children test_remove(avl_tree, 5) # Remove a node with 1 child test_remove(avl_tree, 4) # Remove a node with 2 children result_node = avl_tree.search(7) print("\nFound node object: {}, node value = {}".format(result_node, result_node.val))