BST Iterator Design in DSA Rust

Mental Model

We want to look at the numbers in a tree one by one in order, without storing them all at once.

Analogy: Imagine reading a book page by page instead of opening all pages at once. The iterator helps you see the next page only when you ask for it.

    4
   / \
  2   6
 / \ / \
1  3 5  7
↑stack holds left path nodes

Dry Run Walkthrough

Input: BST: 4 with children 2 and 6; 2 has children 1 and 3; 6 has children 5 and 7; iterate next 4 times

Goal: Return the next smallest number in the BST each time next() is called

Step 1: Initialize iterator by pushing left path from root (4) to stack

stack: [4, 2, 1] (top is 1)

Why: We start with the smallest element by going left as far as possible

Step 2: Call next(): pop 1 from stack, 1 has no right child

stack: [4, 2]

Why: 1 is the smallest, so we return it and move on

Step 3: Call next(): pop 2 from stack, push left path of 2's right child (3)

stack: [4, 3]

Why: After 2, next smallest is 3, so we prepare stack for it

Step 4: Call next(): pop 3 from stack, 3 has no right child

stack: [4]

Why: 3 is next smallest, no right child to add

Step 5: Call next(): pop 4 from stack, push left path of 4's right child (6 then 5)

stack: [6, 5] (top is 5)

Why: 4 is next smallest, now stack ready for 5

Result:

stack: [6, 5]
Returned values in order: 1, 2, 3, 4

Annotated Code

DSA Rust

use std::rc::Rc;
use std::cell::RefCell;

// Definition for a binary tree node.
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }
}

struct BSTIterator {
    stack: Vec<Rc<RefCell<TreeNode>>>,
}

impl BSTIterator {
    fn new(root: Option<Rc<RefCell<TreeNode>>>) -> Self {
        let mut iter = BSTIterator { stack: Vec::new() };
        iter.push_left(root);
        iter
    }

    // Push all left children of node to stack
    fn push_left(&mut self, mut node: Option<Rc<RefCell<TreeNode>>>) {
        while let Some(n) = node {
            self.stack.push(Rc::clone(&n)); // push current node
            node = n.borrow().left.clone(); // move to left child
        }
    }

    fn next(&mut self) -> i32 {
        let node = self.stack.pop().expect("next called on empty iterator");
        let val = node.borrow().val;
        // If right child exists, push its left path
        self.push_left(node.borrow().right.clone());
        val
    }

    fn has_next(&self) -> bool {
        !self.stack.is_empty()
    }
}

fn main() {
    // Build the tree:
    //       4
    //      / \
    //     2   6
    //    / \ / \
    //   1  3 5  7

    let n1 = Rc::new(RefCell::new(TreeNode::new(1)));
    let n3 = Rc::new(RefCell::new(TreeNode::new(3)));
    let n5 = Rc::new(RefCell::new(TreeNode::new(5)));
    let n7 = Rc::new(RefCell::new(TreeNode::new(7)));

    let n2 = Rc::new(RefCell::new(TreeNode::new(2)));
    n2.borrow_mut().left = Some(Rc::clone(&n1));
    n2.borrow_mut().right = Some(Rc::clone(&n3));

    let n6 = Rc::new(RefCell::new(TreeNode::new(6)));
    n6.borrow_mut().left = Some(Rc::clone(&n5));
    n6.borrow_mut().right = Some(Rc::clone(&n7));

    let n4 = Rc::new(RefCell::new(TreeNode::new(4)));
    n4.borrow_mut().left = Some(Rc::clone(&n2));
    n4.borrow_mut().right = Some(Rc::clone(&n6));

    let mut iter = BSTIterator::new(Some(Rc::clone(&n4)));

    let mut result = Vec::new();
    for _ in 0..4 {
        if iter.has_next() {
            result.push(iter.next());
        }
    }
    println!("{}", result.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(", "));
}

iter.push_left(root);

initialize stack with leftmost path from root

while let Some(n) = node { self.stack.push(Rc::clone(&n)); node = n.borrow().left.clone(); }

push all left children to stack to reach smallest node

let node = self.stack.pop().expect("next called on empty iterator");

pop top node from stack as next smallest

self.push_left(node.borrow().right.clone());

push left path of right child to stack for next calls

Output

1, 2, 3, 4

Complexity Analysis

Time: O(1) average per next() call because each node is pushed and popped once

Space: O(h) where h is tree height, due to stack holding left path nodes

vs Alternative: Compared to storing all nodes in a list (O(n) space), this uses less memory and processes nodes lazily

Edge Cases

Empty tree

Iterator stack is empty, has_next() returns false, next() panics if called

DSA Rust

let node = self.stack.pop().expect("next called on empty iterator");

Tree with single node

Stack contains one node, next() returns its value, then stack is empty

DSA Rust

while let Some(n) = node { self.stack.push(Rc::clone(&n)); node = n.borrow().left.clone(); }

Right skewed tree

push_left pushes only one node at a time, next() moves to right child

DSA Rust

self.push_left(node.borrow().right.clone());

When to Use This Pattern

When you need to iterate over a BST in sorted order without extra space, use the BST Iterator pattern with a stack to simulate in-order traversal lazily.

Common Mistakes

Mistake: Pushing all nodes into a list at once instead of using a stack for lazy traversal
Fix: Use a stack to push only left children and process nodes on demand

Mistake: Not pushing left path of right child after popping a node
Fix: After popping, always push left path of the popped node's right child

Summary

It lets you get the next smallest value from a BST one by one using a stack.

Use it when you want sorted access to BST elements without storing them all at once.

The key is to keep a stack of left children so the next smallest is always on top.

Practice

1. What is the main purpose of a BST Iterator in the context of a Binary Search Tree?

easy

A. To traverse the BST in ascending order without recursion

B. To insert new nodes into the BST

C. To delete nodes from the BST

D. To balance the BST automatically

Solution

1. Step 1: Understand BST Iterator functionality

   A BST Iterator is designed to traverse the tree in sorted order, typically ascending, without using recursion.

2. Step 2: Identify the main purpose

   It provides a way to get the next smallest element efficiently, which matches To traverse the BST in ascending order without recursion.

3. Final Answer:

   To traverse the BST in ascending order without recursion -> Option A

4. Quick Check:

   BST Iterator = Ascending traversal [OK]

Hint: BST Iterator returns next smallest element in order [OK]

Common Mistakes:

• Confusing iterator with insertion or deletion
• Thinking it balances the tree
• Assuming it uses recursion internally

2. Which of the following Rust code snippets correctly declares a struct for a BST Iterator holding a stack of nodes?

easy

A. struct BSTIterator { stack: Vec<&TreeNode> }

B. struct BSTIterator { stack: Vec<TreeNode> }

C. struct BSTIterator { stack: Vec> }

D. struct BSTIterator { stack: Vec }

Solution

1. Step 1: Understand ownership in Rust for tree nodes

   Tree nodes are usually boxed to allow recursive types. The iterator needs ownership or mutable access, so using Vec is correct.

2. Step 2: Analyze options

   struct BSTIterator { stack: Vec } uses Vec, which is correct. struct BSTIterator { stack: Vec<TreeNode> } uses Vec<TreeNode> which is invalid due to recursive type size. struct BSTIterator { stack: Vec<&TreeNode> } uses references which complicate lifetimes. struct BSTIterator { stack: Vec> } uses Option unnecessarily.

3. Final Answer:

   struct BSTIterator { stack: Vec } -> Option D

4. Quick Check:

   Use Box for recursive nodes in Vec [OK]

Hint: Use Box to store recursive tree nodes in Vec [OK]

Common Mistakes:

• Using Vec<TreeNode> causes size errors
• Using references without lifetimes
• Wrapping nodes in Option unnecessarily

3. Given the following Rust code for a BST Iterator, what will be printed after calling next() twice?

let mut iter = BSTIterator::new(root);
println!("{}", iter.next().unwrap());
println!("{}", iter.next().unwrap());

Assuming the BST contains nodes with values 3, 7, 9 in sorted order.

medium

A. 9
7

B. 7
9

C. 3
7

D. 3
9

Solution

1. Step 1: Understand BST Iterator traversal order

   The iterator returns nodes in ascending order: 3, then 7, then 9.

2. Step 2: Trace the two next() calls

   First next() returns 3, second next() returns 7.

3. Final Answer:

   3 7 -> Option C

4. Quick Check:

   Next returns ascending values [OK]

Hint: next() returns smallest remaining node [OK]

Common Mistakes:

• Assuming next() returns largest first
• Mixing up order of traversal
• Ignoring unwrap() usage

4. Consider this Rust code snippet for BST Iterator's next() method:

fn next(&mut self) -> Option<i32> {
    let node = self.stack.pop()?;
    let val = node.val;
    let mut current = node.right;
    while let Some(n) = current {
        self.stack.push(n);
        current = n.left;
    }
    Some(val)
}

What is the bug in this code?

medium

A. It pushes nodes in wrong order causing incorrect traversal

B. It does not handle empty stack correctly

C. It moves current incorrectly causing infinite loop

D. It pushes nodes without boxing causing ownership errors

Solution

1. Step 1: Analyze node pushing order in while loop

   The code pushes nodes while moving left, but it pushes current node before going left, which reverses order and breaks in-order traversal.

2. Step 2: Identify correct pushing order

   Nodes should be pushed starting from leftmost node upwards, so pushing before moving left causes wrong order.

3. Final Answer:

   It pushes nodes in wrong order causing incorrect traversal -> Option A

4. Quick Check:

   Push left nodes in correct order [OK]

Hint: Push left children before moving left to maintain order [OK]

Common Mistakes:

• Ignoring order of stack pushes
• Assuming pop handles empty stack safely
• Confusing ownership with traversal logic

5. You want to design a BST Iterator in Rust that supports peek() to see the next smallest element without advancing the iterator. Which approach is best to implement peek() efficiently?

hard

A. Traverse the entire BST and store all elements in a vector upfront

B. Store the next smallest element in a variable updated on each next() call

C. Use recursion inside peek() to find the next element every time

D. Modify the BST structure to add parent pointers for easy access

Solution

1. Step 1: Understand peek() requirement

   Peek should return next smallest element without moving iterator forward, so it needs to know next element in advance.

2. Step 2: Evaluate approaches

   Store the next smallest element in a variable updated on each next() call keeps next element cached and updates it on next(), making peek() O(1). Traverse the entire BST and store all elements in a vector upfront uses extra memory and upfront cost. Use recursion inside peek() to find the next element every time is inefficient due to repeated traversal. Modify the BST structure to add parent pointers for easy access complicates BST unnecessarily.

3. Final Answer:

   Store the next smallest element in a variable updated on each next() call -> Option B

4. Quick Check:

   Cache next element for O(1) peek [OK]

Hint: Cache next element to peek without advancing [OK]

Common Mistakes:

• Recomputing next element on every peek
• Storing entire traversal upfront wasting memory
• Changing BST structure unnecessarily