BST Iterator Design in DSA Go

Mental Model

A BST iterator lets you visit nodes in order one by one without storing all nodes at once. It uses a stack to remember where to go next.

Analogy: Imagine reading a book with bookmarks. You keep bookmarks on pages you might return to later, so you can continue reading in order without flipping all pages at once.

    4
   / \
  2   6
 / \ / \
1  3 5  7

Stack holds nodes to visit next, like bookmarks:
Top -> [1, 2, 4] means next visit is 1, then 2, then 4.

Dry Run Walkthrough

Input: BST: 4,2,6,1,3,5,7; iterate next 4 times

Goal: Visit the first 4 nodes in ascending order: 1, 2, 3, 4

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

Stack: [4, 2, 1]
Next node to visit is top of stack (1)

Why: We push left children to reach smallest node first

Step 2: Call next(): pop 1 from stack, return 1, push right child (none)

Stack: [4, 2]
Returned: 1

Why: 1 has no right child, so just return it

Step 3: Call next(): pop 2 from stack, return 2, push right child (3) and its left path

Stack: [4, 3]
Returned: 2

Why: After visiting 2, we visit its right subtree starting at 3

Step 4: Call next(): pop 3 from stack, return 3, push right child (none)

Stack: [4]
Returned: 3

Why: 3 has no right child, so just return it

Step 5: Call next(): pop 4 from stack, return 4, push right child (6) and its left path (5)

Stack: [6, 5]
Returned: 4

Why: After 4, visit right subtree starting at 6

Result:

Stack: [6, 5]
Returned sequence: 1, 2, 3, 4

Annotated Code

DSA Go

package main

import "fmt"

// TreeNode defines a node in BST
type TreeNode struct {
    Val   int
    Left  *TreeNode
    Right *TreeNode
}

// BSTIterator struct holds stack for controlled traversal
type BSTIterator struct {
    stack []*TreeNode
}

// Constructor initializes iterator and pushes left path
func Constructor(root *TreeNode) BSTIterator {
    it := BSTIterator{stack: []*TreeNode{}}
    it.pushLeft(root)
    return it
}

// pushLeft pushes all left children of node to stack
func (it *BSTIterator) pushLeft(node *TreeNode) {
    for node != nil {
        it.stack = append(it.stack, node) // push node
        node = node.Left                 // go left
    }
}

// Next returns next smallest element
func (it *BSTIterator) Next() int {
    n := len(it.stack)
    node := it.stack[n-1]             // peek top
    it.stack = it.stack[:n-1]         // pop top
    it.pushLeft(node.Right)           // push left path of right child
    return node.Val                   // return current node value
}

// HasNext returns true if iterator has more nodes
func (it *BSTIterator) HasNext() bool {
    return len(it.stack) > 0
}

func main() {
    // Build BST
    root := &TreeNode{4, nil, nil}
    root.Left = &TreeNode{2, &TreeNode{1, nil, nil}, &TreeNode{3, nil, nil}}
    root.Right = &TreeNode{6, &TreeNode{5, nil, nil}, &TreeNode{7, nil, nil}}

    it := Constructor(root)
    for i := 0; i < 4; i++ {
        if it.HasNext() {
            fmt.Print(it.Next())
            if i < 3 {
                fmt.Print(", ")
            }
        }
    }
    fmt.Println()
}

it.pushLeft(root)

push left path from root to stack to start with smallest node

node := it.stack[n-1] it.stack = it.stack[:n-1]

pop top node from stack as next smallest

it.pushLeft(node.Right)

push left path of right child to continue in-order traversal

return node.Val

return current node value as next element

return len(it.stack) > 0

check if more nodes remain to visit

Output

1, 2, 3, 4

Complexity Analysis

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

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

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

Edge Cases

Empty BST (root is nil)

Iterator stack is empty, HasNext returns false immediately

DSA Go

func Constructor(root *TreeNode) BSTIterator {
    it.pushLeft(root) // handles nil root gracefully
}

BST with single node

Stack contains one node, Next returns it, then HasNext returns false

DSA Go

func (it *BSTIterator) HasNext() bool {
    return len(it.stack) > 0
}

BST with only left children

Stack initially contains all nodes, Next pops them in order

DSA Go

func (it *BSTIterator) pushLeft(node *TreeNode) {
    for node != nil {
        it.stack = append(it.stack, node)
        node = node.Left
    }
}

When to Use This Pattern

When you need to traverse a BST in sorted order but want to do it step-by-step without extra memory for all nodes, use BST Iterator pattern with a stack to simulate recursion.

Common Mistakes

Mistake: Pushing right child directly without pushing its left descendants
Fix: Always push left path of right child to stack to maintain in-order order

Mistake: Not checking if stack is empty before calling Next()
Fix: Use HasNext() to check availability before Next()

Summary

Provides a way to iterate BST nodes in ascending order using a controlled stack.

Use when you want lazy, memory-efficient in-order traversal of BST.

The key is to push left children to stack so next smallest is always on top.