The Complete TreeMap

Here is the whole class in outline: its fields and the signature of every method, grouped by where each one comes from. The full code is in the chapter’s repository.

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> implements Map.Entry<K, V> {
    K key;
    V value;
    Node<K, V> left;
    Node<K, V> right;
    int height;                                   // cached height of this subtree
  }

  public TreeMap()

  // The Map operations.
  public V get(K key)                             // search, then read the value
  public boolean containsKey(K key)               // search
  public boolean put(K key, V value)              // add or replace, then rebalance
  public boolean remove(K key)                    // remove, then rebalance
  public int size()
  public Iterator<K> iterator()                   // keys in ascending order
  public Iterable<Map.Entry<K, V>> entries()      // pairs in ascending key order

  private Node<K, V> find(Node<K, V> node, K key)  // contains, returning the node instead of a boolean

  // From the search tree, compared by key. put and remove also refresh and rebalance.
  private Node<K, V> put(Node<K, V> node, K key, V value)
  private Node<K, V> remove(Node<K, V> node, K key)
  private Node<K, V> minNode(Node<K, V> node)     // leftmost node, for the two-child case

  // Balancing.
  private int height(Node<K, V> node)             // cached field, -1 for null
  private int balanceFactor(Node<K, V> node)      // height(left) - height(right)
  private void updateHeight(Node<K, V> node)      // from the children's cached heights
  private Node<K, V> rebalance(Node<K, V> node)   // pick and apply the rotation, if any
  private Node<K, V> rightRotation(Node<K, V> root)
  private Node<K, V> leftRotation(Node<K, V> root)
  private Node<K, V> leftRightRotation(Node<K, V> root)   // left on the child, then right on the node
  private Node<K, V> rightLeftRotation(Node<K, V> root)   // right on the child, then left on the node

  // Iteration: the search tree's in-order iterator, plus the map chapter's Entries.
  private class EntryIterator implements Iterator<Map.Entry<K, V>>
  private class KeyIterator implements Iterator<K>
  private class Entries implements Iterable<Map.Entry<K, V>>
}