Iterating in Key Order

The Map interface also asks for two traversals: iterator(), which visits the keys, and entries(), which visits the key–value pairs. The search tree chapter already wrote an iterator for a binary search tree. It does an in-order traversal with an explicit stack, so it can pause between calls to next.

We use that iterator to visit the nodes, and return each node itself instead of its value. A node implements Map.Entry, so this traversal is exactly what entries() needs:

private class EntryIterator implements Iterator<Map.Entry<K, V>> {
  private Stack<Node<K, V>> stack = new LinkedStack<>();

  EntryIterator() {
    pushLeftSpine(root);
  }

  // Pushes node and every node on the path of left children below it.
  private void pushLeftSpine(Node<K, V> node) {
    while (node != null) {
      stack.push(node);
      node = node.left;
    }
  }

  @Override
  public boolean hasNext() {
    return !stack.isEmpty();
  }

  @Override
  public Map.Entry<K, V> next() {
    if (!hasNext()) {
      throw new NoSuchElementException();
    }
    Node<K, V> node = stack.top();
    stack.pop();
    pushLeftSpine(node.right);
    return node;
  }
}

entries() returns an Iterable, the same way it did in the map chapter: a small class whose iterator() starts a new traversal each time it is called.

@Override
public Iterable<Map.Entry<K, V>> entries() {
  return new Entries();
}

private class Entries implements Iterable<Map.Entry<K, V>> {
  @Override
  public Iterator<Map.Entry<K, V>> iterator() {
    return new EntryIterator();
  }
}

The key iterator runs an EntryIterator and returns the key of each pair, so the traversal is written only once:

@Override
public Iterator<K> iterator() {
  return new KeyIterator();
}

private class KeyIterator implements Iterator<K> {
  private Iterator<Map.Entry<K, V>> nodes = new EntryIterator();

  @Override
  public boolean hasNext() {
    return nodes.hasNext();
  }

  @Override
  public K next() {
    return nodes.next().key();
  }
}

When the traversal is over, nodes.next() throws NoSuchElementException. So KeyIterator.next meets the iterator contract without its own hasNext check.

An in-order traversal of a binary search tree visits its values from smallest to largest, so a TreeMap visits its keys in ascending order. The Map contract does not promise any order, and LinkedMap and ArrayMap do not visit their keys in sorted order. A rotation changes the tree’s shape but keeps the keys in order, so after any sequence of put and remove calls, the traversal still comes out sorted.

The cost is the same as for the search tree iterator. A full traversal visits each node once, so each call to next is amortized $O (1)$ . The stack holds at most one node per level, which is $O (lo g n)$ on an AVL tree.