Binary Search Tree Iterator — InterviewPrepKit

Problem

Design an iterator that walks a binary search tree in in-order (ascending) order. Implement the class BSTIterator:

BSTIterator(root) — initialize with the root of the BST. The iterator starts positioned before the smallest element.
next() -> int — advance to the next element and return its value.
hasNext() -> bool — return whether an element remains to the right of the cursor.

All calls to next() are guaranteed to be valid (an element exists).

Follow-up: make next() and hasNext() run in average O(1) time using only O(h) memory, where h is the height of the tree.

Examples

Example 1

Input:  ["BSTIterator", "next", "next", "hasNext", "next", "hasNext"]
        [[[7, 3, 15, null, null, 9, 20]], [], [], [], [], []]
Output: [null, 3, 7, true, 9, true]

In-order of the tree is 3, 7, 9, 15, 20 — the iterator yields it one value at a time; hasNext is true while values remain.

Example 2

Tree: [2,1,3]; calls: next(), next(), next(), hasNext()
Output: 1, 2, 3, false

After the last element is consumed, hasNext() turns false.

Constraints

The number of nodes is in the range [1, 10^5].
0 <= Node.val <= 10^6
Up to 10^5 total calls to next and hasNext.

The O(h)-memory follow-up is what makes this Medium: flattening the whole tree up front is O(n) memory.

Think about it first

Hint 1

The easy version: what traversal produces a BST's values in sorted order, and what could you precompute in the constructor?

Hint 2

For O(h) memory, think about how *iterative* in-order traversal works with an explicit stack. What does the stack contain at any moment, and what does popping one node correspond to?

Hint 3

Keep a stack holding the path of "left-edge" nodes: pushing a node and all its left descendants costs amortized O(1) per `next`. Popping gives the next value; then push the popped node's right child's left spine. `hasNext` is just "stack non-empty".

TL;DR

Paused in-order traversal via an explicit left-spine stack — amortized O(1) next/hasNext, O(h) space.

Approach 1 — Naive design: flatten everything up front

This is a design problem with no classic brute force. The naive design runs a full in-order traversal in the constructor, stores the sorted values in a list, and serves them by index.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

from typing import Optional, List


class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.values: List[int] = []
        self._inorder(root)
        self.i = 0

    def _inorder(self, node: Optional[TreeNode]) -> None:
        if node is None:
            return
        self._inorder(node.left)
        self.values.append(node.val)
        self._inorder(node.right)

    def next(self) -> int:
        val = self.values[self.i]
        self.i += 1
        return val

    def hasNext(self) -> bool:
        return self.i < len(self.values)

Complexity: constructor O(n) time; next/hasNext true O(1); but O(n) memory, which is exactly what the follow-up forbids. It also does all the work up front even if the caller reads only a few values.

Approach 2 — Controlled in-order traversal with a left-spine stack

Iterative in-order traversal (left subtree, node, right subtree — the order that visits a BST’s values in sorted order) keeps a stack of exactly the ancestors whose values are still pending. If the iterator owns that stack as instance state, the traversal can be paused after each yielded value and resumed on the next call. Invariant: the stack holds the current node’s left-spine ancestors, so the top is always the smallest unvisited value.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

from typing import Optional, List


class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.stack: List[TreeNode] = []
        self._push_left_spine(root)

    def _push_left_spine(self, node: Optional[TreeNode]) -> None:
        while node is not None:
            self.stack.append(node)
            node = node.left

    def next(self) -> int:
        node = self.stack.pop()
        self._push_left_spine(node.right)
        return node.val

    def hasNext(self) -> bool:
        return len(self.stack) > 0

Walkthrough on Example 1 (tree [7,3,15,null,null,9,20], in-order 3, 7, 9, 15, 20):

graph TD
    A7[7] --> A3[3]
    A7 --> A15[15]
    A15 --> A9[9]
    A15 --> A20[20]

Constructor pushes the left spine of 7: stack = [7, 3].
next(): pop 3 (no right child, nothing pushed) → returns 3; stack = [7].
next(): pop 7, push right child 15 and its left spine: stack = [15, 9] → returns 7.
hasNext(): stack non-empty → true.
next(): pop 9 (leaf) → returns 9; stack = [15].
hasNext(): true — 15 and 20 remain.

Complexity: every node is pushed exactly once and popped exactly once across the iterator’s lifetime, so n calls to next cost O(n) total — amortized O(1) per call (a single call can cost O(h) when it descends a long spine, but those costs are prepaid by cheap calls). Space is O(h): the stack never holds more than one root-to-leaf path. hasNext is true O(1).

Approach 3 — Python generator (same idea, language-level)

Python generators are pausable traversals: yield freezes the recursion, and the interpreter keeps the equivalent of the spine stack for you. This is the idiomatic Python phrasing, though an interviewer will usually still want Approach 2’s explicit mechanics.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

from typing import Optional, Iterator


class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.gen: Iterator[int] = self._inorder(root)
        self.peeked: Optional[int] = None
        self._advance()

    def _inorder(self, node: Optional[TreeNode]) -> Iterator[int]:
        if node is not None:
            yield from self._inorder(node.left)
            yield node.val
            yield from self._inorder(node.right)

    def _advance(self) -> None:
        self.peeked = next(self.gen, None)

    def next(self) -> int:
        val = self.peeked
        self._advance()
        assert val is not None
        return val

    def hasNext(self) -> bool:
        return self.peeked is not None

Walkthrough on Example 2 (tree [2,1,3]): the generator is primed to 1. next() returns 1 and primes 2; next() returns 2 and primes 3; next() returns 3 and priming yields None; hasNext() → False.

Complexity: amortized O(1) per call, O(h) space in generator frames. Caveat: yield from recursion makes each yield cost O(depth) in CPython, so the explicit stack is asymptotically cleaner.

Common pitfalls

Pushing only the node, not its left spine: after popping a node, its right child’s entire left spine must be pushed; pushing just the right child yields values out of order.
Claiming worst-case O(1) next: a single next can walk O(h) edges — the guarantee is amortized O(1), justified by push-once/pop-once.
hasNext that mutates: peeking must not consume; in the stack version hasNext should only inspect len(stack).
Forgetting the constructor’s position contract: the iterator starts before the minimum, so the constructor pushes the spine but must not pop anything.

Pattern takeaway

An iterator over a recursive structure is a paused traversal: whatever the recursive version keeps on the call stack becomes explicit instance state, and each next runs the traversal just far enough to produce one element. The left-spine stack invariant (“stack top = smallest pending value”) is the reusable idea — it powers BST iterators, “kth smallest in a BST”, merging two BSTs, and any lazy in-order consumption under O(h) memory.