Skip to Content
Course content

401: Implementing a Balanced Tree Concept (AVL)

Click on the "Edit" button in the top corner of the screen to edit your slide content.

If you've already built a standard Binary Search Tree (BST), you know the dream: $O(\log n)$ search, insertion, and deletion. But as you've probably noticed, that dream dies the moment you insert sorted data. If you feed a BST the numbers 1, 2, 3, 4, 5 in order, you haven't built a tree; you've built a very expensive linked list. That's where the AVL tree comes in. It's a self-balancing BST that ensures the height difference between left and right subtrees is never more than one.

I've seen a lot of engineers struggle with the implementation of the "rotation" logic. Usually, it's not the rotation itself that fails, but the bookkeeping around it. Take a look at this snippet of a rotate_right implementation a student sent me recently:

def rotate_right(self, y):
    x = y.left
    T2 = x.right

    # Perform rotation
    x.right = y
    y.left = T2

    # Return new root
    return x

The Stale Height Trap

At first glance, this looks mathematically correct. We've shifted the nodes, and the BST property is preserved. However, if you run this in a full AVL implementation, your tree will eventually stop balancing or, worse, throw an error. Why? Because the height attribute of the nodes is now lying to you.

In an AVL tree, every node needs to know its height to calculate the balance factor. In the code above, y was the parent and x was the child. After the rotation, y is now a child of x. The height of y has likely decreased, and the height of x has likely increased. By simply returning x without updating the height attributes, the rest of the tree still thinks y is at its old height. The balance factor calculations for every node above this subtree are now based on "stale" data.

Recalculating Heights Post-Rotation

To fix this, you have to update the heights immediately after the pointers change, and you must do it in the correct order: the child first, then the parent. Here is the corrected version:

def rotate_right(self, y):
    x = y.left
    T2 = x.right

    # Perform rotation
    x.right = y
    y.left = T2

    # Fix: Update heights! 
    # y is now lower in the tree, so update it first.
    y.height = 1 + max(self.get_height(y.left), self.get_height(y.right))
    # x is now the root of this subtree, update it last.
    x.height = 1 + max(self.get_height(x.left), self.get_height(x.right))

    return x

I always recommend a helper method like get_height() that handles None nodes by returning 0 (or -1 depending on your convention), otherwise, your code will be littered with if node is not None checks.

Handling the Double-Rotation Headache

The real trick with AVL trees is realizing that a single rotation doesn't solve every imbalance. You'll encounter four scenarios. The "Left-Left" (LL) and "Right-Right" (RR) cases are straightforward—they require one single rotation. But then you hit the "Left-Right" (LR) and "Right-Left" (RL) cases.

If a node is imbalanced because its left child is too heavy, but that child's right subtree is the culprit, a single right rotation won't fix it; it'll just move the imbalance around. You have to perform a "double rotation." For an LR case, you first perform a rotate_left on the left child, transforming the problem into a standard LL case, and then you rotate_right on the original imbalanced node.

It feels like a lot of gymnastics, but the logic is consistent: you are essentially "straightening" the line before you "pivot" it. If you find yourself getting confused, draw it out on paper. I still do this when I'm implementing complex data structures—trying to visualize pointer swaps in your head is a recipe for a AttributeError: 'NoneType' object has no attribute 'left'.

Maintaining the O(log n) Guarantee

The beauty of this overhead is that it pays for itself. By spending a few extra operations during insertion and deletion to keep the tree balanced, you guarantee that the height of the tree remains logarithmic relative to the number of nodes. This means that even with millions of entries, your search time stays incredibly fast, regardless of the order in which the data was inserted.




📋 Practical Task

Exercise: Building an AVL-Backed Unique ID Registry

You are tasked with building a system that stores unique User IDs. Because the IDs are often generated in sequential blocks, a standard BST would become skewed. Implement a full AVLTree class in Python that maintains balance during insertions.

Requirements:

  • Implement a Node class with key, left, right, and height attributes.
  • Implement a insert(key) method that:
    • Performs a standard BST insertion.
    • Updates the height of every ancestor node.
    • Calculates the balance factor.
    • Performs the necessary rotations (Single or Double) to restore balance.
  • Implement a get_height(node) helper to handle None cases safely.
  • Implement a print_tree() method (using an in-order traversal or a simple recursive indented print) to verify that the keys are sorted and the tree is balanced.

Test Case: Insert the following keys in this exact order: [10, 20, 30, 40, 50, 25]. After these insertions, your root node should be 30, and the tree should be balanced. If your root is 50 or 10, your rotation logic is missing!

Rating
0 0

There are no comments for now.

to be the first to leave a comment.