Java AVL insertion: restoring a height invariant with rotations

Last updated: 29 Sept 20264 min read

tutorial

An AVL tree is a binary search tree whose left and right subtree heights differ by at most one at every node.

Java 8+. This is a complete program using JDK classes.

Track the information a repair needs

A shipment-ID index that receives sorted IDs makes an ordinary unbalanced search tree grow into a chain. This implementation stores a height at each node and repairs the first imbalances on the path back from an insertion.

A rotation changes local child links while preserving in-order key order. A left-right insertion first rotates the left child left, then rotates the parent right. Heights are recomputed from the repaired children; updating a parent before its old child is repaired can keep a stale height.

This is a set of int identifiers. Re-inserting an existing identifier does not create another node. A map variant would need a defined replacement policy for its value. For ordinary application code, a library sorted map or set is usually a smaller maintenance commitment than a custom tree.

Working program

public class ShipmentAvlIndex {
    static final class Node {
        final int id; Node left, right; int height = 1;
        Node(int id) { this.id = id; }
    }
    Node root;
    static int height(Node n) { return n == null ? 0 : n.height; }
    static void update(Node n) { n.height = 1 + Math.max(height(n.left), height(n.right)); }
    static Node right(Node n) {
        Node replacement = n.left; n.left = replacement.right; replacement.right = n;
        update(n); update(replacement); return replacement;
    }
    static Node left(Node n) {
        Node replacement = n.right; n.right = replacement.left; replacement.left = n;
        update(n); update(replacement); return replacement;
    }
    static Node insert(Node n, int id) {
        if (n == null) return new Node(id);
        if (id < n.id) n.left = insert(n.left, id);
        else if (id > n.id) n.right = insert(n.right, id);
        else return n;
        update(n);
        int balance = height(n.left) - height(n.right);
        if (balance > 1) {
            if (id > n.left.id) n.left = left(n.left);
            return right(n);
        }
        if (balance < -1) {
            if (id < n.right.id) n.right = right(n.right);
            return left(n);
        }
        return n;
    }
    void add(int id) { root = insert(root, id); }
    boolean contains(int id) {
        Node current = root;
        while (current != null) {
            if (id == current.id) return true;
            current = id < current.id ? current.left : current.right;
        }
        return false;
    }
    public static void main(String[] args) {
        ShipmentAvlIndex index = new ShipmentAvlIndex();
        for (int id : new int[]{30, 10, 20, 40, 50}) index.add(id);
        System.out.println(index.contains(20));
        System.out.println(index.contains(21));
        System.out.println(height(index.root));
    }
}

Output

true
false
3

Costs and boundaries

The maintained height invariant gives O(log n) search and insertion. Insertion uses O(log n) recursion stack; the tree stores O(n) nodes. This implementation covers insertion and lookup, not deletion. Deletion needs its own rebalance cases and tests; the unbalanced BST deletion lesson cannot simply be pasted into this class.

Common Mistakes

Update both affected heights after each rotation.
Do not claim AVL deletion is implemented here.
Test ascending, descending, duplicate and zig-zag insertions.

Continue the contract

Deletion can reduce several ancestor heights and needs its own rotation selection, including a heavy child with zero balance. The deletion companion checks the complete surviving path.

Open the companion lesson.

java

avl-tree