Solution: Caching Subtree Sizes

Both queries are easy if every node knows how many nodes its subtree holds. So we cache that number in the node, next to the height:

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;
  int size;  // number of nodes in the subtree rooted here

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

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

Here is the example tree with each node drawn as key:size:

                 ╭──────╮
                 │ 50:9 │
                 ╰───┬──╯
         ┌───────────┴───────────┐
     ╭───┴──╮                ╭───┴──╮
     │ 30:3 │                │ 70:5 │
     ╰───┬──╯                ╰───┬──╯
    ┌────┴────┐          ┌───────┴───────┐
╭───┴──╮  ╭───┴──╮   ╭───┴──╮        ╭───┴──╮
│ 20:1 │  │ 40:1 │   │ 60:1 │        │ 80:3 │
╰──────╯  ╰──────╯   ╰──────╯        ╰───┬──╯
                                    ┌────┴────┐
                                ╭───┴──╮  ╭───┴──╮
                                │ 75:1 │  │ 90:1 │
                                ╰──────╯  ╰──────╯

The two helpers work the same way as the height helpers. size returns $0$ for the empty tree, and updateSize rebuilds a node’s size from its children’s:

// Returns the cached size of node. The empty tree has size 0.
private int size(Node<K, V> node) {
  if (node == null) {
    return 0;
  }
  return node.size;
}

// Assumes node is not null and its children's cached sizes are correct.
private void updateSize(Node<K, V> node) {
  node.size = 1 + size(node.left) + size(node.right);
}

Like a height, a size is built from the children’s, so it has to be maintained bottom-up, children before parent. The code already calls updateHeight at every node whose subtree might have changed: each node on the path back up from a put or remove, and the two nodes a rotation moves. Those are also the only nodes whose size might change. So updateSize goes in exactly the places updateHeight already is. In the put and remove helpers, it goes right after updateHeight on the way back up:

  updateHeight(node);
  updateSize(node);
  return rebalance(node);

In each single rotation, it goes after the two height updates, in the same order: the node that moved down first, then the node that moved up, whose size depends on it:

private Node<K, V> rightRotation(Node<K, V> root) {
  Node<K, V> child = root.left;
  root.left = child.right;
  child.right = root;
  updateHeight(root);  // root is now below child
  updateHeight(child);
  updateSize(root);
  updateSize(child);
  return child;
}

leftRotation gets the same two lines. The double rotations need nothing, because they are two single rotations. The places that skip updateHeight skip updateSize for the same reason. The new leaf’s constructor sets its size to $1$ . The overwrite branch of put adds no node, so no size changes. The leaf and one-child branches of remove return a subtree that did not change.

The map’s size

The root’s size is the number of nodes in the whole tree. That is what size() returns, so the map no longer needs its own size field:

@Override
public int size() {
  return size(root);
}

The size++ in the public put and the size-- in the public remove go away too. put still checks containsKey first, because it still has to return whether the key was new.

Cost

updateSize is $O (1)$ , and it runs once per node on the path back up and twice per rotation. So put and remove stay $O (lo g n)$ . Each node stores one more int.