put, Final Version

The final version of the helper calls updateHeight on each node before calling rebalance:

private Node<K, V> put(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 = put(node.left, key, value);
  } else if (cmp > 0) {
    node.right = put(node.right, key, value);
  } else {
    node.value = value;  // no new node, so the shape did not change
    return node;
  }
  updateHeight(node);
  return rebalance(node);
}

Each call updates its node after the calls below it have updated their nodes, so every height is computed from correct children. The new leaf needs no update, because its constructor sets its height to $0$ .

Calling updateHeight before rebalance is a choice, not a requirement. rebalance reads only the heights of the node’s children and grandchildren, and a rotation recomputes the height of every node it moves. Calling updateHeight first keeps the last line short: whether or not rebalance rotates, the node it returns has a correct height.

Why put calls updateHeight

The rotations already call updateHeight, and put still needs its own call. A rotation happens only at a node that is out of balance. The other nodes on the path do not rotate, but they can still get taller. Put 4 into a tree that holds 2, 1, and 3:

   2
┌──┴──┐
1     3
      └──┐
         4

Every node is still balanced, so nothing rotates. But the height of 3 went from $0$ to $1$ , and the height of 2 went from $1$ to $2$ . Without the call in put, both nodes would keep their old heights, and the next put would compute wrong balance factors from them. At a node that does rotate, the call in put is redundant, because the rotation computes the same heights again.

The descent visits $O (h)$ nodes. On the way back up, each of those nodes costs one updateHeight and one rebalance, both $O (1)$ now that a height is a field read. So put is $O (h)$ , and because the tree stays balanced, $O (h) = O (lo g n)$ .

In the AVL tree chapter, we saw that an insertion needs at most one repair: the repair brings the subtree back to the height it had before the insertion, so nothing above it can be out of balance. We could stop at the first repair. The code does not stop there. It keeps checking every node on the way up. Above the repair, updateHeight and rebalance change nothing, and they cost $O (1)$ each, so the bound stays $O (lo g n)$ . Checking every node also keeps the code simple, because every node gets the same two steps: refresh its height, then rebalance it.