Implementing remove

The public remove has the same shape as put: it checks for the key first, so it can return the right boolean and keep size correct:

@Override
public boolean remove(K key) {
  if (key == null) {
    throw new IllegalArgumentException();
  }
  if (!containsKey(key)) {
    return false;
  }
  root = remove(root, key);
  size--;
  return true;
}

The recursive helper is the search tree’s remove, compared by key, with the same two lines at the end that put has: updateHeight, then rebalance.

// Removes key from the subtree rooted at node and returns that subtree's
// new root. Assumes key is in the subtree.
private Node<K, V> remove(Node<K, V> node, K key) {
  int cmp = key.compareTo(node.key);
  if (cmp < 0) {
    node.left = remove(node.left, key);
  } else if (cmp > 0) {
    node.right = remove(node.right, key);
  } else if (node.left == null) {
    return node.right;
  } else if (node.right == null) {
    return node.left;
  } else {
    Node<K, V> successor = minNode(node.right);
    node.key = successor.key;
    node.value = successor.value;
    node.right = remove(node.right, successor.key);
  }
  updateHeight(node);
  return rebalance(node);
}

The search tree’s version starts with a base case for null, so removing a missing value does nothing. Here the public method has already checked that the key is there, so the helper can assume it and needs no null case.

Once the key is found, the leaf, one-child, and two-child cases are the search tree’s, with one change. The search tree copies only the successor’s value into the node. A map node holds a pair, so here we copy both the successor’s key and its value, then remove the successor from the right subtree. The minNode helper, which follows left children down to the smallest key, carries over unchanged.

The leaf and one-child branches return right away, because the node they were called on is the one being removed, and the subtree they return was not changed. Every node that stays falls through to the two lines at the bottom. Those are the nodes on the path from the root down to the node that is actually unlinked. In the two-child case, that path goes through the node that took the successor’s key and on down to the successor.

In the AVL tree chapter, we saw that a removal can need repairs at several ancestors, because a repair after a removal can leave the subtree shorter than it was. The code needs no extra logic for that. It follows the same rule as put: refresh the height, then rebalance, at every node on the way up. That fixes each ancestor that is out of balance as the recursion returns through it. And remove is $O (lo g n)$ , the same as put.