Binary Search Tree Iterator | BST | O(H) Space

Designing a Binary Search Tree (BST) Iterator requires creating a stateful iterator class that yields elements in ascending order (in-order traversal order) efficiently, adhering to strict time and space complexity constraints. This chapter covers the step-by-step evolution from a naive storage approach to an optimal $O(H)$ space solution utilizing an explicit stack and controlled left-node traversal.

1. Problem Statement & Functional Requirements

When given the root of a Binary Search Tree, you must design a BSTIterator class that supports two core operations:

• next(): Returns the next smallest element in the BST.
• has_next(): Returns a boolean indicating whether there are still elements remaining in the traversal.

BST Example and In-Order Sequence

Consider a BST with root 7:

        7
       / \
      3   15
     /   /  \
    2   9   20
   /
  1

The corresponding In-Order Traversal (Left -> Node -> Right) yields: [1, 2, 3, 7, 9, 15, 20]

The iterator must expose methods that allow sequential access through this sorted array without computing or storing the entire traversal upfront.

2. Naive Approach: Storing the Full In-Order Traversal

The most intuitive first instinct is to pre-compute the entire in-order traversal during initialization and store it in a linear data structure (like an array or vector).

Naive Implementation Strategy

• Constructor: Perform a complete recursive or iterative in-order traversal of the BST and store every node's value in a vector<int>. Maintain an internal pointer starting at index 0.
• next(): Return the element at the current pointer index and increment the pointer.
• has_next(): Return true if the pointer is less than the total size of the vector; otherwise, return false.

Complexity Analysis

• Time Complexity:
  • Constructor: $O(N)$ to traverse all nodes.
  • next(): $O(1)$ time.
  • has_next(): $O(1)$ time.
• Space Complexity: $O(N)$ because all $N$ node values are permanently stored in memory.

Note: While $O(1)$ time per call is achieved, storing $O(N)$ space violates the interview and system design requirement of operating in $O(H)$ space (where $H$ is the height of the tree).

3. Core Mental Model: Simulating In-Order Traversal with a Stack

To optimize space down to $O(H)$, we must abandon pre-computing the entire array. Instead, we generate elements lazily on-demand.

Real-Life Analogy: The Lazy Box of Books

Imagine stacking books vertically by height. You can only look at the top book. If someone asks for the next smallest book, you pop the top one. If that book has a sub-stack of books on its right, you unpack them and place their leftmost items onto your stack. You never unpack the whole library at once—only what is needed to reach the next item, maintaining a maximum stack height proportional to the tallest shelf ($H$).

In an in-order traversal, the sequence is always Left -> Node -> Right.

1. We dive as deep as possible into the Left subtree, pushing nodes onto a stack.
2. The top of the stack is always the current smallest element (Node).
3. When we pop a node, we check if it has a Right child. If it does, we shift to that right child and dive deep into its left subtree.

4. Step-by-Step Algorithm & Flow Walkthrough

We utilize an explicit Stack data structure to track our path through the BST.

Constructor Initialization (BSTIterator(TreeNode* root))

1. Create an empty stack of tree nodes.
2. Invoke a helper function push_all(root) which pushes the root and all its left descendants onto the stack (root, root->left, root->left->left, ...).

next() Execution Flow

1. Fetch Top: The top of the stack holds the next smallest element in the BST.
2. Pop: Remove the top node from the stack and store its reference.
3. Check Right Subtree: If the popped node has a right child (node->right), call push_all(node->right). This dives into the right subtree's extreme left chain and pushes those nodes onto the stack.
4. Return: Return the value of the popped node.

has_next() Execution Flow

• Simply check if the stack is non-empty (!stack.empty()). If nodes remain on the stack, return true; otherwise, return false.

5. Complexity Breakdown: Why $O(1)$ Amortized Time?

At first glance, pushing an entire left chain inside next() might look like an $O(H)$ operation per call. However, let's analyze the aggregate cost across all operations:

Thus, across $N$ total next() calls, the total push operations equal $N$. Dividing $N$ pushes across $N$ calls yields an amortized $O(1)$ time complexity per call. Space complexity remains strictly $O(H)$ because the stack never holds more nodes than the height of the tree at any given instant.

6. Production-Ready Code Implementation

Below are the complete, production-grade implementations in C++ and Java.

C++ Implementation

#include <stack>

class BSTIterator {
private:
    std::stack<TreeNode*> st;

    // Helper function to push all left descendants onto the stack
    void pushAll(TreeNode* node) {
        for (; node != nullptr; st.push(node), node = node->left);
    }

public:
    // Constructor initializes the iterator with the root of the BST
    BSTIterator(TreeNode* root) {
        pushAll(root);
    }
    
    // Returns the next smallest element
    int next() {
        TreeNode* tmpNode = st.top();
        st.pop();
        
        // If the popped node has a right child, process its left descendants
        if (tmpNode->right != nullptr) {
            pushAll(tmpNode->right);
        }
        
        return tmpNode->val;
    }
    
    // Returns whether we have a next element
    bool hasNext() {
        return !st.empty();
    }
};

Java Implementation

import java.util.Stack;

class BSTIterator {
    private Stack<TreeNode> stack;

    // Constructor initializes the iterator with the root of the BST
    public BSTIterator(TreeNode root) {
        stack = new Stack<>();
        pushAll(root);
    }
    
    // Returns the next smallest element
    public int next() {
        TreeNode tmpNode = stack.pop();
        
        // If the popped node has a right child, process its left descendants
        if (tmpNode.right != null) {
            pushAll(tmpNode.right);
        }
        
        return tmpNode.val;
    }
    
    // Returns whether we have a next element
    public boolean hasNext() {
        return !stack.isEmpty();
    }

    // Helper method to push all left descendants onto the stack
    private void pushAll(TreeNode node) {
        while (node != null) {
            stack.push(node);
            node = node.left;
        }
    }
}

7. Summary and Key Takeaways

• Problem Goal: Design a memory-efficient iterator for a Binary Search Tree supporting next() and has_next().
• Space Limitation: Standard storage of the entire in-order traversal requires $O(N)$ space, which fails optimal constraints.
• Optimal Space ($O(H)$): Achieved by simulating an iterative in-order traversal using an explicit stack that tracks only the current branch height.
• Lazy Evaluation: Nodes are pushed and popped dynamically only when requested, avoiding upfront full-tree traversal overhead.
• Amortized $O(1)$ Time: Although a single next() call may trigger multiple pushes down a right subtree's left chain, every node is pushed and popped exactly once, guaranteeing $O(1)$ time on average per call.
• Clean Encapsulation: Utilizing a dedicated private helper (pushAll) keeps constructor and next() logic clean, modular, and easy to maintain.