Starting from the Search Tree

An AVL tree is a binary search tree, so we start from the search tree chapter’s BinarySearchTree class and change what the map needs.

The first change is the node. The search tree’s node held one value. A map’s node holds a key and the value stored under it. And like LinkedMap’s node, it implements Map.Entry, so entries() can return the nodes themselves as read-only pairs:

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;

    Node(K key, V value) {
      this.key = key;
      this.value = value;
    }

    @Override
    public K key() {
      return key;
    }

    @Override
    public V value() {
      return value;
    }
  }

  public TreeMap() {
    root = null;
    size = 0;
  }

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

  // ... the operations will go here
}

The second change is the size field. The search tree had no size, but Map has a size() method, and keeping a count makes it $O(1)$. As with LinkedMap, the rule is that size always equals the number of nodes, so every operation that adds or removes a node updates it.