Update avl_tree.py

Author: lighting9999

data_structures/binary_tree/avl_tree.py

@@ -1,17 +1,8 @@
-"""
-Auto-balanced binary tree implementation
-For doctests: python3 -m doctest -v avl_tree.py
-For testing: python avl_tree.py
-"""
-
 from __future__ import annotations
 
-import math
 import random
-import doctest
 from typing import Any
 
-
 class MyQueue:
     def __init__(self) -> None:
         self.data: list[Any] = []
@@ -25,204 +16,176 @@ def push(self, data: Any) -> None:
         self.tail += 1
 
     def pop(self) -> Any:
-        ret = self.data[self.head]
-        self.head += 1
-        return ret
-
+        return self.data[self.head]
 
 class MyNode:
+    __slots__ = ("data", "left", "right", "height")
+    
     def __init__(self, data: Any) -> None:
         self.data = data
-        self.left = self.right = None
+        self.left: MyNode | None = None
+        self.right: MyNode | None = None
         self.height = 1
 
-    def get_data(self) -> Any:
-        return self.data
-
-    def get_left(self) -> MyNode | None:
-        return self.left
-
-    def get_right(self) -> MyNode | None:
-        return self.right
-
-    def get_height(self) -> int:
-        return self.height
-
-    def set_data(self, data: Any) -> None:
-        self.data = data
-
-    def set_left(self, node: MyNode | None) -> None:
-        self.left = node
-
-    def set_right(self, node: MyNode | None) -> None:
-        self.right = node
-
-    def set_height(self, height: int) -> None:
-        self.height = height
-
-
-def get_height(node: MyNode | None) -> int:
+def get_height(node: MyNode | None) -> int: 
     return node.height if node else 0
 
-
-def my_max(a: int, b: int) -> int:
+def my_max(a: int, b: int) -> int: 
     return a if a > b else b
 
-
 def right_rotation(node: MyNode) -> MyNode:
-    print("left rotation node:", node.data)
-    ret = node.left
-    node.left = ret.right
-    ret.right = node
+    left_child = node.left
+    if left_child is None:
+        return node
+    
+    node.left = left_child.right
+    left_child.right = node
     node.height = my_max(get_height(node.right), get_height(node.left)) + 1
-    ret.height = my_max(get_height(ret.right), get_height(ret.left)) + 1
-    return ret
-
+    left_child.height = my_max(get_height(left_child.right), get_height(left_child.left)) + 1
+    return left_child
 
 def left_rotation(node: MyNode) -> MyNode:
-    print("right rotation node:", node.data)
-    ret = node.right
-    node.right = ret.left
-    ret.left = node
+    right_child = node.right
+    if right_child is None:
+        return node
+    
+    node.right = right_child.left
+    right_child.left = node
     node.height = my_max(get_height(node.right), get_height(node.left)) + 1
-    ret.height = my_max(get_height(ret.right), get_height(ret.left)) + 1
-    return ret
-
+    right_child.height = my_max(get_height(right_child.right), get_height(right_child.left)) + 1
+    return right_child
 
 def lr_rotation(node: MyNode) -> MyNode:
-    node.left = left_rotation(node.left)
+    if node.left:
+        node.left = left_rotation(node.left)
     return right_rotation(node)
 
-
 def rl_rotation(node: MyNode) -> MyNode:
-    node.right = right_rotation(node.right)
+    if node.right:
+        node.right = right_rotation(node.right)
     return left_rotation(node)
 
-
 def insert_node(node: MyNode | None, data: Any) -> MyNode | None:
-    if not node:
+    if node is None:
         return MyNode(data)
-
+    
     if data < node.data:
         node.left = insert_node(node.left, data)
         if get_height(node.left) - get_height(node.right) == 2:
-            if data < node.left.data:
+            if node.left and data < node.left.data:
                 node = right_rotation(node)
             else:
                 node = lr_rotation(node)
     else:
         node.right = insert_node(node.right, data)
         if get_height(node.right) - get_height(node.left) == 2:
-            if data < node.right.data:
+            if node.right and data < node.right.data:
                 node = rl_rotation(node)
             else:
                 node = left_rotation(node)
-
+    
     node.height = my_max(get_height(node.right), get_height(node.left)) + 1
     return node
 
-
-def get_extreme(root: MyNode, is_right: bool) -> Any:
-    while child := root.right if is_right else root.left:
-        root = child
+def get_left_most(root: MyNode) -> Any:
+    while root.left:
+        root = root.left
     return root.data
 
-
-def del_node(root: MyNode, data: Any) -> MyNode | None:
-    if root.data == data:
+def del_node(root: MyNode | None, data: Any) -> MyNode | None:
+    if root is None:
+        return None
+        
+    if data == root.data:
         if root.left and root.right:
-            root.data = get_extreme(root.right, False)
+            root.data = get_left_most(root.right)
             root.right = del_node(root.right, root.data)
         else:
             return root.left or root.right
-    elif root.data > data:
-        if not root.left:
-            return root
+    elif data < root.data:
         root.left = del_node(root.left, data)
     else:
         root.right = del_node(root.right, data)
-
-    # Handle balancing
-    right_height = get_height(root.right)
+    
+    if root.left is None and root.right is None:
+        root.height = 1
+        return root
+    
     left_height = get_height(root.left)
-
+    right_height = get_height(root.right)
+    
     if right_height - left_height == 2:
-        if get_height(root.right.right) > get_height(root.right.left):
+        right_right = get_height(root.right.right) if root.right else 0
+        right_left = get_height(root.right.left) if root.right else 0
+        if right_right > right_left:
             root = left_rotation(root)
         else:
             root = rl_rotation(root)
-    elif right_height - left_height == -2:
-        if get_height(root.left.left) > get_height(root.left.right):
+    elif left_height - right_height == 2:
+        left_left = get_height(root.left.left) if root.left else 0
+        left_right = get_height(root.left.right) if root.left else 0
+        if left_left > left_right:
             root = right_rotation(root)
         else:
             root = lr_rotation(root)
-
+    
     root.height = my_max(get_height(root.right), get_height(root.left)) + 1
     return root
-
-
 class AVLtree:
-    def __init__(self) -> None:
-        self.root = None
-
-    def get_height(self) -> int:
+    __slots__ = ("root",)
+    
+    def __init__(self) -> None: 
+        self.root: MyNode | None = None
+        
+    def get_height(self) -> int: 
         return get_height(self.root)
-
+    
     def insert(self, data: Any) -> None:
-        print(f"insert:{data}")
         self.root = insert_node(self.root, data)
-
+    
     def del_node(self, data: Any) -> None:
-        print(f"delete:{data}")
-        if not self.root:
-            return
         self.root = del_node(self.root, data)
-
+    
     def __str__(self) -> str:
-        if not self.root:
+        if self.root is None:
             return ""
-        q, output, layer, cnt = MyQueue(), "", self.get_height(), 0
-        q.push(self.root)
-
-        while not q.is_empty():
-            node = q.pop()
-            space = " " * int(2 ** (layer - 1))
-            output += space + (str(node.data) if node else "*") + space
-            cnt += 1
-
-            if node:
-                q.push(node.left)
-                q.push(node.right)
+            
+        levels = []
+        queue = [self.root]
+        
+        while queue:
+            current = []
+            next_level = []
+            
+            for node in queue:
+                if node:
+                    current.append(str(node.data))
+                    next_level.append(node.left)
+                    next_level.append(node.right)
+                else:
+                    current.append("*")
+                    next_level.append(None)
+                    next_level.append(None)
+            
+            if any(node is not None for node in next_level):
+                levels.append(" ".join(current))
+                queue = next_level
             else:
-                q.push(None)
-                q.push(None)
-
-            for i in range(10):
-                if cnt == 2**i - 1:
-                    layer -= 1
-                    output += "\n"
-                    if layer == 0:
-                        break
-                    break
-
-        return output + "\n" + "*" * 36
-
-
-def _test() -> None:
-    doctest.testmod()
-
+                break
+        
+        return "\n".join(levels) + "\n" + "*"*36
 
-if __name__ == "__main__":
-    _test()
+def test_avl_tree() -> None:
     t = AVLtree()
     lst = list(range(10))
     random.shuffle(lst)
-
+    
     for i in lst:
         t.insert(i)
-        print(t)
-
+    
     random.shuffle(lst)
     for i in lst:
         t.del_node(i)
-        print(t)
+
+if __name__ == "__main__":
+    test_avl_tree()