DSA A BST Iterator

The problem

A BST iterator is an object that hands out the values of a binary search tree in sorted order, one per call. The method next returns the smallest value not yet returned. The method has_next reports whether any value remains.

Calls can mix in any order. A caller may check has_next between any two next calls, or not at all. The operation script below is the test every design in this article must pass.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6, TreeNode(4))),
                TreeNode(10, None, TreeNode(14)))

script = ["has_next", "next", "next", "has_next", "next", "next",
          "next", "next", "next", "has_next"]
expected = [True, 1, 3, True, 4, 6, 8, 10, 14, False]
for step, want in zip(script, expected):
    print("call", step, "->", want)

The tree holds seven values, from 1 to 14. The script drains it fully. The final check answers False, because nothing remains after 14.

Intuition: the values come out in inorder

Sorted order on a BST is inorder order. An inorder walk visits the left subtree, then the node, then the right subtree. The BST rule makes that walk read values from smallest to largest.

An iterator is that walk, paused between values. Each next call resumes the walk for exactly one value. The design problem is how to store a paused walk.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6, TreeNode(4))),
                TreeNode(10, None, TreeNode(14)))

order = []
def inorder(node):
    if node is None:
        return
    inorder(node.left)
    order.append(node.value)
    inorder(node.right)
inorder(tree)
print("inorder:", order)
print("inorder is sorted:", order == sorted(order))
print("calls would hand out:", order[0], "then", order[1], "then", order[2])

The sorted check is the evidence. The calls hand out the inorder list from its start. Any storage scheme that reproduces this list, one value at a time, satisfies the spec.

Brute force: the naive design collects everything first

The naive design runs the full inorder walk inside its constructor. It stores every value in a list. A separate index counts how many values have been handed out.

Each call is then trivial. Next reads one list slot and moves the index. Has_next compares the index with the list length.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

class ListIterator:
    def __init__(self, root):
        self.values = []
        self.index = 0
        def inorder(node):
            if node is None:
                return
            inorder(node.left)
            self.values.append(node.value)
            inorder(node.right)
        inorder(root)

    def has_next(self):
        return self.index < len(self.values)

    def next(self):
        value = self.values[self.index]
        self.index += 1
        return value

tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6, TreeNode(4))),
                TreeNode(10, None, TreeNode(14)))

iterator = ListIterator(tree)
print("values collected before the first call:", len(iterator.values))
results = []
for step in ["has_next", "next", "next", "has_next", "next", "next",
             "next", "next", "next", "has_next"]:
    if step == "has_next":
        results.append(iterator.has_next())
    else:
        results.append(iterator.next())
print("script results:", results)
print("results match the spec:", results == [True, 1, 3, True, 4, 6, 8, 10, 14, False])

The script passes in full. The cost sits in the first line of the output. All seven values are collected and stored before the first call, even if the caller only ever takes two.

The optimal solution: a stack holds the path to the next value

A stack inside the iterator holds the path from the root down to the smallest value not yet returned, and each next call pops that value and pushes the path into its right subtree. That one sentence is the whole design.

A stack is a list used at one end, where append adds an item and pop removes the most recent one. The top of this stack is always the next value waiting. Has_next is then one length check.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

class BSTIterator:
    def __init__(self, root):
        self.stack = []
        self.pushes = 0
        self.pops = 0
        self.push_path(root)

    def push_path(self, node):
        while node is not None:
            self.stack.append(node)
            self.pushes += 1
            node = node.left

    def has_next(self):
        return len(self.stack) > 0

    def next(self):
        node = self.stack.pop()
        self.pops += 1
        self.push_path(node.right)
        return node.value

tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6, TreeNode(4))),
                TreeNode(10, None, TreeNode(14)))

iterator = BSTIterator(tree)
results = []
for step in ["has_next", "next", "next", "has_next", "next", "next",
             "next", "next", "next", "has_next"]:
    if step == "has_next":
        results.append(iterator.has_next())
    else:
        results.append(iterator.next())
print("script results:", results)
print("results match the spec:", results == [True, 1, 3, True, 4, 6, 8, 10, 14, False])
print("total pushes:", iterator.pushes, "total pops:", iterator.pops)

The script passes with the same results as the naive design. The counters close the proof. Seven nodes produce exactly seven pushes and seven pops across the whole drain.

Walkthrough: the stack through a full drain

The walkthrough drains the tree and prints the stack after every call. Stack values print bottom first. The top of the stack, printed last, is always the next value waiting.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

class BSTIterator:
    def __init__(self, root):
        self.stack = []
        self.pushes = 0
        self.pops = 0
        self.push_path(root)

    def push_path(self, node):
        while node is not None:
            self.stack.append(node)
            self.pushes += 1
            node = node.left

    def has_next(self):
        return len(self.stack) > 0

    def next(self):
        node = self.stack.pop()
        self.pops += 1
        self.push_path(node.right)
        return node.value

tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6, TreeNode(4))),
                TreeNode(10, None, TreeNode(14)))

iterator = BSTIterator(tree)
print("after creation, stack:", [node.value for node in iterator.stack])
while iterator.has_next():
    value = iterator.next()
    print("next returned", value, "- stack now:",
          [node.value for node in iterator.stack])
print("has_next at the end:", iterator.has_next())

Creation pushes 8, 3 and 1, the path to the smallest value. Returning 3 pushes the path through 6 down to 4. Returning 8 pushes 10, whose right child 14 follows later.

The stack empties exactly when the last value leaves. Has_next then answers False. The paused walk was stored in that one stack the whole time.

Complexity

The naive design takes O(n) time and O(n) space before its first call, for n nodes. The stack design takes O(h) time before its first call, for height h, because it pushes only one root-to-leaf path.

Amortised cost is the total cost of many calls spread evenly over those calls. Each node enters the stack once and leaves once, so a full drain costs O(n) across n calls, which is O(1) amortised per call. The stack holds at most one root-to-leaf path, so its space is O(h).

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

class BSTIterator:
    def __init__(self, root):
        self.stack = []
        self.pushes = 0
        self.pops = 0
        self.push_path(root)

    def push_path(self, node):
        while node is not None:
            self.stack.append(node)
            self.pushes += 1
            node = node.left

    def has_next(self):
        return len(self.stack) > 0

    def next(self):
        node = self.stack.pop()
        self.pops += 1
        self.push_path(node.right)
        return node.value

def complete_tree(values):
    if not values:
        return None
    mid = len(values) // 2
    return TreeNode(values[mid], complete_tree(values[:mid]),
                    complete_tree(values[mid + 1:]))

def max_stack(root):
    iterator = BSTIterator(root)
    biggest = len(iterator.stack)
    while iterator.has_next():
        iterator.next()
        if len(iterator.stack) > biggest:
            biggest = len(iterator.stack)
    return biggest, iterator.pushes

for n in [7, 15, 31]:
    tree = complete_tree(list(range(1, n + 1)))
    biggest, pushes = max_stack(tree)
    print("nodes", n, "- naive list stored:", n,
          " stack design max stack:", biggest, " total pushes:", pushes)

The counts prove both claims. The list grows with the node count itself. The largest stack grows only when the tree gains a level, while total pushes still equal the node count.

Edge cases and pitfalls

An empty tree makes an iterator that is finished at birth. Its stack is empty and has_next answers False at once. A single node answers one value and then finishes.

Repeated has_next calls are harmless. They change no state, as the run below shows. The next call after three checks still returns the smallest value.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

class BSTIterator:
    def __init__(self, root):
        self.stack = []
        self.pushes = 0
        self.pops = 0
        self.push_path(root)

    def push_path(self, node):
        while node is not None:
            self.stack.append(node)
            self.pushes += 1
            node = node.left

    def has_next(self):
        return len(self.stack) > 0

    def next(self):
        node = self.stack.pop()
        self.pops += 1
        self.push_path(node.right)
        return node.value

empty = BSTIterator(None)
print("empty tree has_next:", empty.has_next(), "stack size:", len(empty.stack))
single = BSTIterator(TreeNode(9))
print("single node has_next:", single.has_next(), "next:", single.next(),
      "has_next after:", single.has_next())

tree = TreeNode(8,
                TreeNode(3, TreeNode(1), TreeNode(6, TreeNode(4))),
                TreeNode(10, None, TreeNode(14)))

iterator = BSTIterator(tree)
print("has_next three times:", iterator.has_next(), iterator.has_next(), iterator.has_next())
print("next still returns the smallest:", iterator.next())
left_chain = TreeNode(5, TreeNode(4, TreeNode(3, TreeNode(2, TreeNode(1)))))
chain_iterator = BSTIterator(left_chain)
print("left chain stack at creation:", [node.value for node in chain_iterator.stack])
try:
    empty.next()
    print("next on empty: no error")
except IndexError:
    print("next on empty raises: IndexError")

Two pitfalls close the article. A left chain pushes its whole length at creation, because its one root-to-leaf path is the whole tree, so stack space is O(h) and a chain makes h equal to n. And next on a finished iterator raises IndexError, so callers guard every next call with has_next, exactly as the spec's script does.