Implementing a Self-Balancing AVL Tree in Python

A binary search tree is useful because it keeps smaller keys to the left and larger keys to the right. Its performance depends on the tree’s shape, though. Insert already sorted values such as 10, 20, 30, and 40, and the structure can become a long chain. Searching then takes linear time instead of the expected logarithmic time.

An AVL tree prevents that decline by tracking the height of every subtree and correcting imbalances after updates. The result is a self-balancing binary search tree with predictable performance. This guide builds an AVL tree in Python, explains the four rotations it uses, and shows how to test the implementation. It also fits naturally alongside the data structure and algorithm material on hello ML.

Why AVL trees matter

An AVL tree is named after Georgy Adelson-Velsky and Evgenii Landis. For every node, the heights of its left and right subtrees may differ by at most one. This difference is called the balance factor:

balance factor = height(left subtree) - height(right subtree)

A balance factor of -1, 0, or 1 is valid. A value of 2 means the node is too heavy on the left, while -2 means it is too heavy on the right. The tree repairs itself when an insertion causes either of those values.

Keeping the height difference small means searching, inserting, and deleting remain O(log n) in the worst case. A regular binary search tree can degrade to O(n), particularly when data arrives in sorted or nearly sorted order. That distinction matters in a Python service receiving timestamped events, student records, or sensor readings in a predictable sequence.

AVL trees are stricter than some alternatives, such as red-black trees. They may perform more rotations during updates, but they often provide shorter paths for lookups. For a Melbourne logistics application repeatedly querying indexed locations, or a Sydney team maintaining an in-memory catalogue, that lookup behaviour can be valuable.

Representing nodes and measuring balance

Each node needs a key, references to its children, and its current height. Storing height directly avoids recalculating an entire subtree every time the balance factor is needed. An empty subtree has height zero, and a leaf node has height one.

Here is a compact node definition and the helper code used by the tree:

class Node:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None
        self.height = 1


def height(node):
    return node.height if node is not None else 0


def update_height(node):
    node.height = 1 + max(height(node.left), height(node.right))


def balance_factor(node):
    return height(node.left) - height(node.right)

The height function handles None, which keeps the rotation and insertion code readable. After changing a node’s children, update_height must run before its balance factor is used. A frequent implementation error is to rotate correctly but leave stale heights behind, causing a later insertion to make the wrong decision.

The implementation below treats duplicate keys as no-ops. That policy is simple and appropriate when keys identify unique records. If duplicates matter, a node could hold a counter, a list of values, or a consistent rule that sends equal keys to one side.

Rotations that restore order

A rotation changes the shape of a subtree without changing its in-order key sequence. A right rotation fixes a left-heavy subtree, while a left rotation fixes a right-heavy subtree. The operation updates child references first, then recalculates heights from the lower node upwards.

def rotate_right(y):
    x = y.left
    middle = x.right

    x.right = y
    y.left = middle

    update_height(y)
    update_height(x)
    return x


def rotate_left(x):
    y = x.right
    middle = y.left

    y.left = x
    x.right = middle

    update_height(x)
    update_height(y)
    return y

There are four imbalance patterns. A left-left case needs one right rotation, and a right-right case needs one left rotation. A left-right case first rotates the left child to the left, then rotates the unbalanced node to the right. A right-left case follows the mirror image: rotate the right child right, then rotate the node left.

The key is to inspect both the node’s balance factor and the inserted key. After recursively inserting into the left subtree, a balance factor above one indicates left heaviness. If the new key is greater than the left child’s key, the shape is left-right and needs two rotations. The corresponding comparison on the right identifies a right-left case.

Putting insertion together

Insertion follows ordinary binary search tree logic until the recursive call returns. At that point, each ancestor updates its height and checks its balance factor. If a rotation occurs, the new root of that local subtree is returned to the caller. This return value is important because a rotation can change the root of the complete tree.

class AVLTree:
    def __init__(self):
        self.root = None

    def insert(self, key):
        self.root = self._insert(self.root, key)

    def _insert(self, node, key):
        if node is None:
            return Node(key)

        if key < node.key:
            node.left = self._insert(node.left, key)
        elif key > node.key:
            node.right = self._insert(node.right, key)
        else:
            return node

        update_height(node)
        balance = balance_factor(node)

        # Left-left case
        if balance > 1 and key < node.left.key:
            return rotate_right(node)

        # Right-right case
        if balance < -1 and key > node.right.key:
            return rotate_left(node)

        # Left-right case
        if balance > 1 and key > node.left.key:
            node.left = rotate_left(node.left)
            return rotate_right(node)

        # Right-left case
        if balance < -1 and key < node.right.key:
            node.right = rotate_right(node.right)
            return rotate_left(node)

        return node

The public insert method replaces self.root with the result of the private recursive method. That assignment handles the first rotation at the tree root. Without it, the tree might internally rotate while the external root reference still points to an outdated node.

You can test all four cases with short sequences:

tree = AVLTree()

for value in [30, 20, 10, 25, 28, 40, 50]:
    tree.insert(value)

An in-order traversal should produce sorted output regardless of the rotations:

def inorder(node):
    if node is not None:
        yield from inorder(node.left)
        yield node.key
        yield from inorder(node.right)


print(list(inorder(tree.root)))
# [10, 20, 25, 28, 30, 40, 50]

For Australian developers working through a coding interview or a LeetCode-style exercise, checking sorted traversal is a fast first test. It confirms the search-tree property, while a separate height check confirms the balancing property.

Testing, deletion, and practical use

A robust test suite should verify more than a few hand-picked values. Insert ascending values, descending values, random values, and repeated values. After each operation, confirm that the in-order traversal is sorted and that every node has a balance factor between -1 and 1.

def validate(node):
    if node is None:
        return 0

    left_height = validate(node.left)
    right_height = validate(node.right)

    assert abs(left_height - right_height) <= 1
    assert node.height == 1 + max(left_height, right_height)

    return node.height

Deletion uses the same general pattern as insertion: remove a key using binary search tree rules, update heights while returning from recursion, and rebalance every ancestor. The rotation selection differs slightly because a deletion may produce a child with balance factor zero. It is often best to implement insertion first, add validation, and then extend the same structure for deletion.

AVL operations use O(log n) time when the tree is valid and O(n) space in the worst case for recursive call depth, although balancing keeps that depth logarithmic. Each node uses additional memory for its height and two child references. Python’s object overhead can be significant, so a production workload with millions of records may prefer a database index, a specialised library, or arrays of primitive values.

The right choice depends on workload. An AVL tree is useful when ordered iteration and predictable lookups matter, while a hash table is generally better for direct equality lookup. In Australia, a small Perth startup may favour a clear pure-Python structure for a prototype, whereas a larger Brisbane platform may rely on PostgreSQL indexes or a managed cache for operational simplicity. The implementation remains valuable because it explains what those indexed structures are designed to achieve.

Implementation habits worth keeping

• Store subtree heights on nodes so balance checks remain constant time.
• Return the new subtree root after every recursive insertion and rotation.
• Decide how duplicate keys behave before writing comparison logic.
• Test ascending, descending, random, and repeated input sequences.
• Validate both in-order sorting and the allowed balance-factor range.
• Add deletion only after insertion and rotation cases are well tested.
• Compare the custom tree with built-in dictionaries and database indexes before using it in production.

A useful next step is to compare this implementation with related material in the algorithms collection, then adapt the code to store records with a separate comparison key. Someone saying they will “have a crack” at the exercise during an arvo coding session can start with integer keys, add validation, and only then introduce custom objects.

The central design lesson is that balancing is part of every update, not a repair performed occasionally. Recursive insertion preserves the search-tree ordering, height metadata exposes imbalance, and rotations restore the logarithmic shape. Once those responsibilities are separated into small functions, an AVL tree becomes much easier to understand, test, and extend.