Where Balancing Goes

put can add a new node where there was an empty child, which can make one side a level taller. remove takes a node out, which can make one side a level shorter. Either change can make the heights of a node’s two subtrees differ by more than one, so these are the two operations that have to rebalance.

Before we write put and remove, we need to know where in them the balancing goes. Recall the shape of add from the search tree chapter, written here for keys:

private Node<K, V> add(Node<K, V> node, K key, V value) {
  if (node == null) {
    return new Node<>(key, value);
  }

  int cmp = key.compareTo(node.key);
  if (cmp < 0) {
    node.left = add(node.left, key, value);
  } else if (cmp > 0) {
    node.right = add(node.right, key, value);
  }
  return node;
}

In the AVL tree chapter, we found the nodes to check after an insertion by hand: follow the path back up from the new leaf toward the root, checking each node on the way. Those are the only nodes whose subtrees changed. The recursion already visits exactly those nodes as it unwinds, in that order, and each call returns the root of its subtree for its caller to store. So the balancing goes just before each call returns its node: check that node’s balance, rotate if it is out of balance, and return whatever is now the root of the subtree. Removal is also recursive and also returns the root of each subtree, so the balancing goes in the same place in removal.

For the balancing we need two things: the rotations, and a way to read a node’s balance factor and pick a rotation. We write the rotations first.