Implementing put and remove

The questions in this section use this outline of TreeMap. Assume rebalance(node) repairs node if it is unbalanced and returns the subtree’s new root, and that containsKey and minNode (the node with the smallest key in a subtree) are already written.

public class TreeMap<K extends Comparable<K>, V> implements Map<K, V> {
  private Node<K, V> root;
  private int size;

  private static class Node<K, V> {
    K key;
    V value;
    Node<K, V> left;
    Node<K, V> right;
    int height;  // cached height of this subtree

    // Sets key and value, with no children and height 0.
    Node(K key, V value) { ... }
  }

  // Returns node.height, or -1 if node is null.
  private int height(Node<K, V> node) { ... }

  // Sets node.height from its children's cached heights.
  private void updateHeight(Node<K, V> node) { ... }

  // Returns height(node.left) - height(node.right).
  private int balanceFactor(Node<K, V> node) { ... }
}

Exercise

Write the recursive helper behind put. It adds a new node if key is not in the subtree, replaces the value if it is, keeps the subtree balanced and every cached height correct, and returns the subtree’s new root.

private Node<K, V> put(Node<K, V> node, K key, V value) {

}

Solution

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;
    return node;
  }
  updateHeight(node);
  return rebalance(node);
}

Exercise

The public put throws IllegalArgumentException if the key or the value is null. It returns true if the key was new and false if it replaced a value. Write it so that it also keeps size correct.

public boolean put(K key, V value) {

}

Solution

public boolean put(K key, V value) {
  if (key == null || value == null) {
    throw new IllegalArgumentException();
  }
  boolean isNew = !containsKey(key);
  root = put(root, key, value);
  if (isNew) {
    size++;
  }
  return isNew;
}

Exercise

The recursive helper private Node<K, V> remove(Node<K, V> node, K key) removes key from the subtree rooted at node, which must contain it, and returns the subtree’s new root. Below is the end of that helper. In the last branch, node holds the key being removed and has two children. Write the code for this branch using the in-order successor.

} else {
  // node holds key and has two children

}
updateHeight(node);
return rebalance(node);

Solution

} 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);

Solution

A repair after a removal can leave the subtree shorter than it was before the removal. Its parent can then be unbalanced, and so on up to the root.