DSA — AVL Trees

What is an AVL Tree?

A BST that automatically balances itself after each operation.

class AVLNode:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None
        self.height = 1

Height Balance

def get_height(node):
    if not node:
        return 0
    return node.height

def get_balance(node):
    if not node:
        return 0
    return get_height(node.left) - get_height(node.right)

Rotations

# Right rotation
def right_rotate(y):
    x = y.left
    T2 = x.right
    x.right = y
    y.left = T2
    y.height = 1 + max(get_height(y.left), get_height(y.right))
    x.height = 1 + max(get_height(x.left), get_height(x.right))
    return x

# Left rotation
def left_rotate(x):
    y = x.right
    T2 = y.left
    y.left = x
    x.right = T2
    x.height = 1 + max(get_height(x.left), get_height(x.right))
    y.height = 1 + max(get_height(y.left), get_height(y.right))
    return y

Insert with Balancing

def insert(node, data):
    if not node:
        return AVLNode(data)
    if data < node.data:
        node.left = insert(node.left, data)
    else:
        node.right = insert(node.right, data)
    
    node.height = 1 + max(get_height(node.left), get_height(node.right))
    balance = get_balance(node)
    
    # Left Left
    if balance > 1 and data < node.left.data:
        return right_rotate(node)
    # Right Right
    if balance < -1 and data > node.right.data:
        return left_rotate(node)
    # Left Right
    if balance > 1 and data > node.left.data:
        node.left = left_rotate(node.left)
        return right_rotate(node)
    # Right Left
    if balance < -1 and data < node.right.data:
        node.right = right_rotate(node.right)
        return left_rotate(node)
    
    return node

Mini Practice

1. Implement AVL tree insert
2. Perform rotations
3. Maintain balance
4. Compare with regular BST

Up Next

Continue with Heaps — heap data structure.