Topics › Binary Search Trees › Practice Problems

BST iterator

Hard TUF ID 96 SDE

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

Design an iterator over a binary search tree that returns its keys in ascending order (the in-order sequence). It is constructed with the root of the tree and supports two operations:

• next() returns the next smallest key and advances. It is only called when a key remains.
• hasNext() returns whether any key remains, without consuming anything.

Constraints (LeetCode 173): 1 to 10^5 nodes, keys between 0 and 10^6, at most 10^5 calls in total. The follow-up asks for next() and hasNext() in average O(1) time using O(h) memory, where h is the height of the tree.

The Java file describes all three implementations with one small interface, so a single test script can drive each of them:

interface InorderIterator {
    boolean hasNext();
    int next();
}

Example 1

        7
       / \
      3   15
         /  \
        9    20

calls:   next  next  hasNext  next  hasNext  next  hasNext  next  hasNext
returns: 3     7     true     9     true     15    true     20    false

Example 2

      2
     / \
    1   3

calls:   hasNext  hasNext  next  next  next  hasNext
returns: true     true     1     2     3     false

Calling hasNext() twice in a row must not skip a key.

Approach 1: Brute force, flatten the tree up front

Intuition. The in-order traversal of a BST visits keys in sorted order, because every node is visited after its smaller left subtree and before its larger right subtree. If we run that traversal once in the constructor and save the keys in a list, the iterator becomes a cursor into an array: next() returns keys[pos++] and hasNext() checks pos < size.

This is simple and every operation after construction is truly O(1). The cost is memory: it stores all n keys even if the caller only wants the first two, and it pays the whole O(n) traversal before the first next().

Algorithm.

1. In the constructor, do a recursive in-order traversal, appending each key to keys.
2. hasNext(): return pos < keys.size().
3. next(): return keys.get(pos) and increment pos.

static class FlattenedIterator implements InorderIterator {
    private final List<Integer> keys = new ArrayList<>();
    private int pos = 0;                             // index of the next key to hand out

    FlattenedIterator(TreeNode root) {
        fill(root);
    }

    private void fill(TreeNode node) {
        if (node == null) return;
        fill(node.left);
        keys.add(node.val);
        fill(node.right);
    }

    public boolean hasNext() {
        return pos < keys.size();
    }

    public int next() {
        if (!hasNext()) throw new NoSuchElementException("no more keys");
        return keys.get(pos++);
    }
}

Complexity. Constructor O(n) time. next() and hasNext() O(1). Space O(n) for the list, which fails the follow-up’s O(h) memory requirement.

Approach 2: Optimal, controlled in-order with a stack

Intuition. Look at what the iterative in-order traversal keeps on its stack: exactly the ancestors whose key has not been returned yet, ordered so the smallest one is on top. That is all the state an iterator needs. Instead of running the traversal to completion, we pause it after each visited node and resume it on the next call.

Two facts make this work:

• The smallest key in any subtree is at the end of its leftmost path. So pushing a node and all its left descendants puts the subtree’s minimum on top of the stack.
• The in-order successor of a node x is either the minimum of x.right or, if x has no right child, the nearest ancestor whose left subtree contains x. The first case is handled by pushing the left chain of x.right. The second case is already handled: that ancestor is the next node down in the stack.

So next() pops the top, pushes the left chain of its right child, and returns the popped key. hasNext() is just “is the stack non-empty”.

Algorithm.

1. Constructor: push root and every node on its leftmost path.
2. hasNext(): return !stack.isEmpty().
3. next(): pop node; push the left chain starting at node.right; return node.val.

static class StackIterator implements InorderIterator {
    private final Deque<TreeNode> stack = new ArrayDeque<>();

    StackIterator(TreeNode root) {
        pushLeftChain(root);
    }

    private void pushLeftChain(TreeNode node) {
        while (node != null) {                       // the smallest unvisited key ends up on top
            stack.push(node);
            node = node.left;
        }
    }

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

    public int next() {
        if (stack.isEmpty()) throw new NoSuchElementException("no more keys");
        TreeNode node = stack.pop();
        pushLeftChain(node.right);                   // the successor is the leftmost node of the right subtree
        return node.val;
    }
}

Dry run on Example 1 (stack written bottom to top):

Complexity. hasNext() is O(1). A single next() can push up to h nodes, but every node is pushed exactly once and popped exactly once over the whole iteration. So n calls to next() cost O(n) in total: amortised O(1) per call. Space O(h): the stack never holds more than one root-to-leaf path. Both follow-up requirements are met.

This iterator is also the building block for problems that need a sorted stream from a BST without materialising it, such as Two Sum in a BST with a forward and a backward iterator.

Approach 3: Morris iterator, constant extra space

Intuition. The stack in Approach 2 remembers “where to return after finishing a left subtree”. Morris traversal stores that same information in the tree’s unused right pointers. Before going into cur’s left subtree, find cur’s in-order predecessor (the rightmost node in that left subtree) and point its right at cur. When the walk later reaches the predecessor and follows its right, it arrives back at cur. Finding the thread already there means the left subtree is finished: remove the thread and return cur’s key.

As an iterator, the state is just one pointer cur. next() runs the Morris loop until it reaches a visit, returns that key, and pauses. hasNext() is cur != null. That is correct because after each visit cur moves to a real right child or a thread to the successor, and only becomes null after the very last key, which has neither.

Algorithm.

1. Constructor: cur = root.
2. hasNext(): return cur != null.
3. next(): loop:
   • If cur.left is null: remember cur.val, move to cur.right, return the key.
   • Otherwise find pred. If pred.right is null, set it to cur and go left. If it equals cur, reset it to null, remember cur.val, move right, return the key.

static class MorrisIterator implements InorderIterator {
    private TreeNode cur;                            // the next key comes from cur or from its left subtree

    MorrisIterator(TreeNode root) {
        cur = root;
    }

    public boolean hasNext() {
        return cur != null;
    }

    public int next() {
        while (cur != null) {
            if (cur.left == null) {                  // nothing smaller is pending: cur is next
                int val = cur.val;
                cur = cur.right;                     // a real right child or a thread to the successor
                return val;
            }
            TreeNode pred = cur.left;                // in-order predecessor: rightmost node of the left subtree
            while (pred.right != null && pred.right != cur) pred = pred.right;
            if (pred.right == null) {
                pred.right = cur;                    // first arrival: thread back, then go left
                cur = cur.left;
            } else {
                pred.right = null;                   // second arrival: left subtree is done, restore the tree
                int val = cur.val;
                cur = cur.right;
                return val;
            }
        }
        throw new NoSuchElementException("no more keys");
    }
}

Complexity. O(1) extra space and amortised O(1) per next(): across a full iteration each edge is walked a constant number of times. hasNext() is O(1).

The catch. The tree is modified while the iteration is in progress, and it is only fully restored once the last key has been returned. If a caller abandons the iterator halfway, the tree is left with cycles, and any other reader running at the same time sees a corrupted tree. The Java tests drain every iterator and confirm the tree’s shape is unchanged afterwards, but in production code you would choose the stack iterator unless memory is extremely tight and you control the whole iteration.

Edge cases and pitfalls

• hasNext() must be side-effect free. All three versions only inspect state. A version that advances inside hasNext() skips keys when it is called twice.
• Empty tree. LeetCode guarantees at least one node, but all three versions report hasNext() == false immediately and throw NoSuchElementException from next() instead of returning garbage.
• Skewed trees. A left chain makes the constructor of the stack iterator push all n nodes (h = n), and a right chain makes each next() push one node. Both are still amortised O(1).
• Iterating while modifying. None of these iterators tolerate insertions or deletions during iteration; the stack can hold detached nodes and the flattened list goes stale.
• Recursion depth. Approach 1’s recursive fill can overflow Java’s stack on a chain of 10^5 nodes. The other two are iterative.

Related problems

• Two sum in BST runs an ascending and a descending iterator as two pointers.
• Iterative Inorder Traversal of Binary Tree is the same stack logic as one loop.
• Morris Inorder Traversal of a Binary Tree explains threading in full.
• Find K-th smallest element in BST is “call next() k times”.
• Inorder Successor/Predecessor in BST finds a single successor directly from the root.