Caching the Height

Give Node a height field that holds the height of the subtree rooted there:

private static class Node<K, V> implements Map.Entry<K, V> {
  K key;
  V value;
  Node<K, V> left;
  Node<K, V> right;
  int height;

  Node(K key, V value) {
    this.key = key;
    this.value = value;
    height = 0;  // a new node is a leaf
  }

  // key() and value() as before
}

With the height stored in the node, the height helper reads the field. The empty tree has no node to hold a number, so it still answers $-1$ for null:

private int height(Node<K, V> node) {
  if (node == null) {
    return -1;
  }
  return node.height;
}

Every place that asks for a height still calls height(node), so none of them changes. The cost changes. Getting a height used to be an $O(n)$ traversal, and now it is an $O(1)$ field read. balanceFactor reads two heights, so it becomes $O(1)$ too.

Keeping size correct is an invariant of TreeMap: size always equals the number of nodes. Keeping each node’s height correct is an invariant of the class too: a node’s height field always equals the height of its subtree. Just as we update size when an operation adds or removes an entry, we need to update a node’s height when an operation changes the subtree below it. That happens in put, in remove, and in every rotation. When a node’s two children already hold the right heights, the node’s own height follows from them by the formula that defines height:

private void updateHeight(Node<K, V> node) {
  node.height = 1 + Math.max(height(node.left), height(node.right));
}

This is the same formula the recursive height uses, without the recursive calls. It reads two fields and writes one, so it is $O(1)$. It is only right if the children’s heights are already right, so the cache has to be maintained bottom-up, children before parent.

put and remove already visit nodes in that order. They recurse down to the position where they insert or remove a node and return back up through every ancestor, child before parent. If each call runs updateHeight on its node just before returning it, every node’s height is updated right after its children’s heights.