AVLTree - master - Monsimento/Algoritms

AVLTree
103 строки3 KB
Monsimento
create AVLTree
11 ноя 2025, 09:12
11 ноя 2025, 09:12cad715d
class AVLNode:
    def __init__(self, val):
        self.val = val
        self.left: AVLNode = None
        self.right: AVLNode = None
        self.height = 1

class AVLTree:
    def height(self, node: AVLNode):
        if not node:
            return 0
        return node.height
    
    def balance_factor(self, node: AVLNode):
        if not node:
            return 0
        return 1 + self.height(node.right) - self.height(node.left)
    
    def update_height(self, node: AVLNode):
        if node:
            node.height = 1 + max(self.height(node.left), self.height(node.right))

    def rotate_right(self, z: AVLNode):
        y:AVLNode = z.right
        t3:AVLNode = z.left

        y.right = z
        z.left = t3

        self.update_height(z)
        self.update_height(y)

        return y
    
    def rotate_left(self, y:AVLNode):
        z:AVLNode = y.right
        t3:AVLNode = z.left

        y.right = z
        z.left = t3
        
        self.update_height(y)
        self.update_height(z)

        return z

    def insert(self, node: AVLNode, val):
        # 1 base insert
        if not node:
            return(AVLNode(val))
        if val < node.val:
            node.left = self.insert(node.left, val)
        elif val > node.val:
            node.right = self.insert(node.right, val)
        else:
            return node
        
        # 2 update height

        self.update_height(node)

        # balance_factor 

        bf = self.balance_factor(node) 

        if bf < -1 and val < node.left.val:
            return self.rotate_right(node)
        if bf > -1 and val > node.right.val:
            return self.rotate_left(node)
        
        if bf <-1 and val > node.left.val:
            node.left = self.rotate_right(node.left)
            return self.rotate_right(node)
        if bf >-1 and val < node.right.val:
            node.right = self.rotate_right(node.right)
            return self.rotate_left(node)
        
        return node
    
    def inorder(self, node:AVLNode, res):
        if node:
            self.inorder(node.left, res)
            res.append(node.val)
            self.inorder(node.right, res)

def printer(root, lvl=0, prefix="root:"):
    if root:
        print(" " *(lvl * 4) + prefix + str(root.val))
        if root.left:
            printer(root.left, lvl+1, 'l-----')
        else:
            printer(root.right, lvl+1, 'r-----')


tree = AVLTree()
root = None
for i in range(100):
    root = tree.insert(root, i)

res = []
tree.inorder(root, res)
print(f'inorder {res}')
printer(root)