Iterating in Sorted Order

We want a binary search tree to work in an enhanced for loop, the way our sets and maps did:

for (T value : tree) {
  // ...
}

For that, the tree has to implement Iterable<T>:

public class BinarySearchTree<T extends Comparable<T>>
    implements Iterable<T> {
  // ...
}

Its iterator() method returns a fresh iterator each time it is called:

@Override
public Iterator<T> iterator() {
  return new BinarySearchTreeIterator();
}

An iterator has to traverse the tree. Any of the four traversal orders from the tree chapter would visit every value once. But in-order traversal visits a binary search tree’s values from smallest to largest. We showed that earlier in this chapter, when we found that in-order traversal comes out sorted. So we iterate in-order.

We cannot use the recursive in-order traversal, though. Each call to next has to return one value and stop, and a recursive traversal runs to the end once it starts. In the practice problems of the tree chapter, we implemented all three depth-first traversals iteratively, with an explicit stack instead of recursion. One reason to do that was to be able to pause the traversal between values, which is what an iterator needs to do.

Here is the iterative in-order traversal from those practice problems, which we will turn into the binary search tree’s iterator:

public static <T> List<T> inOrder(Node<T> root) {
  List<T> result = new ArrayList<>();
  Stack<Node<T>> stack = new LinkedStack<>();
  Node<T> current = root;
  while (current != null || !stack.isEmpty()) {
    while (current != null) {   // push the left spine
      stack.push(current);
      current = current.left;
    }
    current = stack.top();      // top has no unvisited left, so it is ready
    stack.pop();
    result.add(current.value);  // record it
    current = current.right;    // step into the right subtree
  }
  return result;
}

Three changes turn it into an iterator that stops after each value and resumes on the next call:

• The stack becomes a field, so it survives between calls.
• “Push the left spine” becomes a helper, pushLeftSpine, because we now call it in two places: once to fill the stack from the root when the iterator is created, and once at the end of each next, on the popped node’s right child.
• The loop body (pop the top, record it, step right) becomes the body of next. The current variable disappears: because each call pushes the whole left spine before returning, no half-pushed spine is ever left between calls, so hasNext only has to check whether the stack is empty.

Here is the result. It is an inner class, so it can read the tree’s root:

private class BinarySearchTreeIterator implements Iterator<T> {
  private Stack<Node<T>> stack = new LinkedStack<>();

  BinarySearchTreeIterator() {
    pushLeftSpine(root);
  }

  private void pushLeftSpine(Node<T> node) {
    while (node != null) {
      stack.push(node);
      node = node.left;
    }
  }

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

  @Override
  public T next() {
    if (!hasNext()) {
      throw new NoSuchElementException();
    }
    Node<T> node = stack.top();
    stack.pop();
    T value = node.value;
    pushLeftSpine(node.right);
    return value;
  }
}

When the stack is empty, next throws NoSuchElementException, the same contract as the iterators in the iterator chapter.