Binary Search Tree Iterator

LeetCode 173: Binary Search Tree Iterator Solution & Explanation

Problem Overview

TL;DR: The optimal solution simulates a recursive in-order traversal using an explicit stack to achieve $O(h)$ memory complexity and amortized $O(1)$ time complexity.

The LeetCode 173 problem asks us to design a Binary Search Tree Iterator class. Given the root of a Binary Search Tree (BST), we need to implement an iterator that returns nodes in ascending order (in-order traversal). The iterator must support two operations: next(), which returns the next smallest number, and hasNext(), which checks if there are remaining numbers.

This is a popular interview question because it tests your understanding of tree traversal internals and your ability to convert recursive logic into an iterative state machine.

Brute Force Approach for Binary Search Tree Iterator

A naive approach to solving the Binary Search Tree Iterator problem is to perform a complete in-order traversal of the BST during the initialization phase and store the results in a linear data structure.

Naive Algorithm Steps

1. Initialize: Create a list or array to store node values.
2. Traverse: In the constructor, run a recursive in-order DFS (Depth-First Search) on the entire tree.
3. Store: Append every visited node's value to the list.
4. Iterate: Maintain a pointer (index) starting at -1.
   • next() increments the index and returns the element at that position.
   • hasNext() checks if the index + 1 is within the list bounds.

Pseudo-Code

Class BSTIterator:
    List values
    Integer index

    Constructor(root):
        values = []
        inorder(root, values)
        index = -1

    Function inorder(node, list):
        if node is null return
        inorder(node.left, list)
        list.add(node.val)
        inorder(node.right, list)

    Function next():
        index++
        return values[index]

Analysis & Limitations

• Time Complexity: $O(N)$ for the constructor, where $N$ is the number of nodes. $O(1)$ for next() and hasNext().
• Space Complexity: $O(N)$ to store all node values.
• Why it fails: While this passes the basic constraints, it fails the Follow up requirement to use $O(h)$ memory, where $h$ is the height of the tree. In a massive tree where we might only need to iterate through the first few elements, pre-processing the entire tree is inefficient and wastes memory.

Algorithm Strategy: Tree Traversal Patterns (DFS & BFS)

We implement the iterator by maintaining the state of a partial Depth-First Search.

1. Data Structure: Use a Stack to store TreeNode pointers. This stack replaces the system call stack used in recursion.
2. Helper Function (pushAllLeft):
   • Takes a node as input.
   • Iterates while the node is not null.
   • Pushes the current node onto the stack.
   • Updates the node to its left child (node = node.left).
   • This sets up the stack such that the top element is always the next smallest value in the subtree.
3. Constructor:
   • Initialize the empty stack.
   • Call pushAllLeft(root) to reach the first (smallest) element of the BST.
4. next() Implementation:
   • Pop the node from the top of the stack. This is the return value.
   • Crucial Step: If the popped node has a right child, call pushAllLeft(node.right). This prepares the stack for the next call by processing the right subtree (and its left descendants).
   • Return the value of the popped node.
5. hasNext() Implementation:
   • Simply check if the stack is non-empty.

Execution Flow

Let's trace the execution with a simple BST: root = 7, root.left = 3, root.right = 15.

1. Initialization (BSTIterator(7)):

   • Call pushAllLeft(7).
   • Push 7. Move to 7.left (3).
   • Push 3. Move to 3.left (null).
   • Stop. Stack State: [7, 3] (Top is 3).

2. Call next():

   • Pop 3 from stack. Stack State: [7].
   • Check 3.right. It is null. Do nothing.
   • Return: 3.

3. Call next():

   • Pop 7 from stack. Stack State: [].
   • Check 7.right. It is 15.
   • Call pushAllLeft(15).
   • Push 15. Move to 15.left (null).
   • Stop. Stack State: [15].
   • Return: 7.

4. Call next():

   • Pop 15 from stack. Stack State: [].
   • Check 15.right. It is null.
   • Return: 15.

5. Call hasNext():

   • Stack is empty. Return false.

Proof of Correctness

The algorithm correctly implements an in-order traversal because it strictly follows the recursive definition: process the left subtree, then the current node, then the right subtree.

• By pushing all left children initially, we ensure the first popped element is the leftmost node (smallest).
• By popping a node, we are "visiting" it.
• By immediately processing the right child of the popped node (and that right child's left descendants), we ensure that after a node $N$ is visited, the next node visited is the smallest node in $N$'s right subtree (or $N$'s ancestor if no right subtree exists), which matches the in-order successor logic.

Pattern Reuse Notes

The Tree DFS - Recursive Inorder Traversal pattern and the technique of using an explicit stack to simulate recursion are applicable to several other problems:

• LeetCode 94: Binary Tree Inorder Traversal - This is the fundamental problem that defines the pattern used in the iterator.
• LeetCode 98: Validate Binary Search Tree - You can validate a BST by ensuring the values produced by this inorder iteration are strictly increasing.
• LeetCode 230: Kth Smallest Element in a BST - Instead of iterating fully, stop after calling next() exactly $k$ times.
• LeetCode 501: Find Mode in Binary Search Tree - Use the iterator logic to traverse sorted values and count frequencies without storing the full tree.

Reference Implementation

C++ Solution for LeetCode 173

1/**
2 * Definition for a binary tree node.
3 * struct TreeNode {
4 *     int val;
5 *     TreeNode *left;
6 *     TreeNode *right;
7 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
8 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
9 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
10 * };
11 */
12class BSTIterator {
13private:
14    stack<TreeNode*> st;
15
16    // Helper to push node and all its left descendants
17    void pushAllLeft(TreeNode* node) {
18        while (node != nullptr) {
19            st.push(node);
20            node = node->left;
21        }
22    }
23
24public:
25    BSTIterator(TreeNode* root) {
26        pushAllLeft(root);
27    }
28    
29    int next() {
30        // The top of the stack is the smallest element currently available
31        TreeNode* topNode = st.top();
32        st.pop();
33        
34        // If the node has a right child, we must process its left subtree
35        if (topNode->right != nullptr) {
36            pushAllLeft(topNode->right);
37        }
38        
39        return topNode->val;
40    }
41    
42    bool hasNext() {
43        return !st.empty();
44    }
45};

Java Solution for LeetCode 173

1/**
2 * Definition for a binary tree node.
3 * public class TreeNode {
4 *     int val;
5 *     TreeNode left;
6 *     TreeNode right;
7 *     TreeNode() {}
8 *     TreeNode(int val) { this.val = val; }
9 *     TreeNode(int val, TreeNode left, TreeNode right) {
10 *         this.val = val;
11 *         this.left = left;
12 *         this.right = right;
13 *     }
14 * }
15 */
16class BSTIterator {
17    private Deque<TreeNode> stack;
18
19    public BSTIterator(TreeNode root) {
20        stack = new ArrayDeque<>();
21        pushAllLeft(root);
22    }
23    
24    // Helper to push node and all its left descendants
25    private void pushAllLeft(TreeNode node) {
26        while (node != null) {
27            stack.push(node);
28            node = node.left;
29        }
30    }
31
32    public int next() {
33        // Pop the smallest element
34        TreeNode topNode = stack.pop();
35        
36        // If there is a right subtree, process its left boundary
37        if (topNode.right != null) {
38            pushAllLeft(topNode.right);
39        }
40        
41        return topNode.val;
42    }
43
44    public boolean hasNext() {
45        return !stack.isEmpty();
46    }
47}

Python Solution for LeetCode 173

1# Definition for a binary tree node.
2# class TreeNode:
3#     def __init__(self, val=0, left=None, right=None):
4#         self.val = val
5#         self.left = left
6#         self.right = right
7
8class BSTIterator:
9
10    def __init__(self, root: Optional[TreeNode]):
11        self.stack = []
12        self._push_all_left(root)
13
14    def _push_all_left(self, node):
15        # Helper to push node and all its left descendants
16        while node:
17            self.stack.append(node)
18            node = node.left
19
20    def next(self) -> int:
21        # Pop the smallest element
22        top_node = self.stack.pop()
23        
24        # If there is a right subtree, process its left boundary
25        if top_node.right:
26            self._push_all_left(top_node.right)
27            
28        return top_node.val
29
30    def hasNext(self) -> bool:
31        return len(self.stack) > 0