BST iterator with a stack

Why this matters

A binary search tree keeps its keys in sorted order when you walk it in the right way. The walk that visits the left subtree, then the node, then the right subtree is called the inorder walk. Standard containers such as std::set and std::map give you the same kind of walk through an iterator. An iterator is an object that hands you the next item each time you ask for one, so you never need the whole sequence in memory at once.

This lesson builds that kind of iterator for a BST by hand. It uses a stack, a container that gives back the last item you put in first. The lesson compares the stack version with the simpler method of copying every key into a list first. Here the walk keeps going from one call to the next.

The problem with a full list

The easy way to get sorted keys is to walk the whole tree, write every key into a std::vector, and then read the vector from the front. That works. It has two costs.

The first cost is time before the first answer. To return the smallest key, the program must visit every node, because it cannot know which key is smallest until it has seen them all.

The second cost is memory. The vector holds one entry for every key, even if the caller only wants the first three keys.

An iterator avoids both costs. It returns the smallest key after visiting only the nodes on the way down to it. Later calls do a small amount of work each. The price is a stack that holds a few node addresses at a time.

The idea: a stack of nodes still to visit

Take this tree. It has nine keys. The lesson uses it throughout.

            8
          /   \
         3     10
        / \      \
       1   6      14
          / \    /
         4   7  13

To get the smallest key, walk left from the root as far as you can. Every node you pass is bigger than the node you reach at the end, so those nodes wait. Put each one on a stack as you pass it. Here the walk goes 8, then 3, then 1. The stack now holds 8, 3 and 1, with 1 on top. Node 1 has no left child, so 1 is the smallest key.

To get the next key, take the top of the stack. That is node 1, and its key is the answer. Pop it. Node 1 has no right child, so nothing new is added.

The next key is 3. Node 3 is the new top. Its key is the answer. Pop it. Node 3 has a right child, 6. Node 6 and everything in its subtree are bigger than 3 and smaller than 8, so they come next. Walk left from 6 and push each node on the way: 6, then 4. Node 4 has no left child, so stop. Now the stack holds 4, 6 and 8 with 4 on top.

Each time you pop a node, you do two things. You return its key. You then push the left branch of its right child, if it has one, because that subtree holds the next keys in order.

Step by step on the example tree

Here is every call to next(). The constructor is the code that runs when the iterator is created. Each row shows the value returned and the stack after the call. The top of the stack is on the left.

The rule for each call fits in three steps:

1. Read the key on top of the stack. That key is the answer.
2. Pop it.
3. Push the left branch of its right child, going down from that child to the bottom of its left side.

Step 3 handles the right subtree. Keys inside that subtree are bigger than the node you just returned. They are also smaller than every node still waiting on the stack, because those waiting nodes are ancestors that are bigger still. So the right subtree is the next part of the sorted order.

The code

Here is the iterator class. The tree is built elsewhere. The iterator only reads it.

class BstIterator {
public:
    explicit BstIterator(const Node* root) { push_left(root); }
 
    bool has_next() const { return !stack_.empty(); }
 
    int next() {
        const Node* node = stack_.top();
        stack_.pop();
        push_left(node->right.get());
        return node->key;
    }
 
    std::size_t stack_size() const { return stack_.size(); }
 
private:
    void push_left(const Node* node) {
        for (; node != nullptr; node = node->left.get()) stack_.push(node);
    }
 
    std::stack<const Node*> stack_;
};

Some terms in this code:

• Node is a struct with an int key and two child pointers, left and right. A pointer holds the memory address of another object. A null pointer holds no address and means there is no child.
• std::unique_ptr<Node> is a pointer that owns the object it points to. When the owner goes away, the object is freed. The tree owns its nodes this way. In this code the iterator never frees a node. Its .get() calls return the plain address inside the unique_ptr.
• const Node* is a pointer to a node that the holder may read but not change. The iterator does not change the tree, so it uses this type.
• std::stack<const Node*> is a stack of those pointers. push adds to the top, top reads the top, and pop removes the top without returning it.
• has_next() says whether there is any key left. An empty stack means the walk is finished.

The constructor calls push_left on the root. That fills the stack with the left branch from the root down. The first call to next() then returns the smallest key, as the trace above showed.

Worked example 1: a small tree with the full walk

The first example builds the nine-key tree and prints every key with next(). The tree comes from inserting the keys in this order: 8, 3, 10, 1, 6, 14, 4, 7, 13. The insert function in the full program adds a key at the right place and does no balancing.

Tree root;
for (int k : {8, 3, 10, 1, 6, 14, 4, 7, 13}) root = insert(std::move(root), k);
BstIterator it(root.get());
std::cout << "example 1:";
while (it.has_next()) std::cout << ' ' << it.next();
std::cout << '\n';

Output:

example 1: 1 3 4 6 7 8 10 13 14

The loop asks has_next() before each call. That check matters. Calling next() on an empty iterator is an error, covered in the mistakes section below.

Worked example 2: merging two trees with two iterators

The second example is harder. Two separate trees hold sorted keys. The goal is one sorted list with the keys from both trees. The method is a merge: keep one iterator for each tree and always take the smaller of the two current keys.

The variables va and vb hold the current key of each iterator. The ha and hb flags say whether that iterator still has a key. A flag becomes false when its iterator runs out, and then that tree is ignored.

Tree a, b;
for (int k : {2, 5, 9}) a = insert(std::move(a), k);
for (int k : {1, 4, 6, 10}) b = insert(std::move(b), k);
BstIterator ia(a.get()), ib(b.get());
bool ha = ia.has_next(), hb = ib.has_next();
int va = ha ? ia.next() : 0;
int vb = hb ? ib.next() : 0;
std::cout << "example 2:";
while (ha || hb) {
    if (!hb || (ha && va <= vb)) {
        std::cout << ' ' << va;
        ha = ia.has_next();
        if (ha) va = ia.next();
    } else {
        std::cout << ' ' << vb;
        hb = ib.has_next();
        if (hb) vb = ib.next();
    }
}
std::cout << '\n';

Output:

example 2: 1 2 4 5 6 9 10

Check the first steps by hand. The first key of a is 2 and the first key of b is 1. The smaller is 1, so 1 is printed and the iterator for b moves to 4. Now 2 and 4 are compared, and 2 is printed. The pattern continues until both trees are empty.

The merge uses two iterators at once. Each one keeps its own stack, so the memory used is the two stacks together, not the two trees copied into lists.

Worked example 3: walking from the largest key down

The third example changes the direction of the walk. The iterator now starts at the largest key and goes down. The change is small. Instead of walking left and pushing left children, walk right and push right children. In next(), the iterator then pushes the right branch of the left child.

class DescendingIterator {
public:
    explicit DescendingIterator(const Node* root) { push_right(root); }
 
    bool has_next() const { return !stack_.empty(); }
 
    int next() {
        const Node* node = stack_.top();
        stack_.pop();
        push_right(node->left.get());
        return node->key;
    }
 
private:
    void push_right(const Node* node) {
        for (; node != nullptr; node = node->right.get()) stack_.push(node);
    }
 
    std::stack<const Node*> stack_;
};

Run it on the same tree from example 1:

DescendingIterator di(root.get());
std::cout << "example 3:";
while (di.has_next()) std::cout << ' ' << di.next();
std::cout << '\n';

Output:

example 3: 14 13 10 8 7 6 4 3 1

The two classes are mirror images. Once you understand one, you can write the other by swapping left and right.

Memory: the stack against a full list

The stack holds only nodes on one path. Each node on the stack is either the next key or an ancestor of it. A path from the root to a leaf has at most as many nodes as the tree has levels. Call the number of nodes on the longest such path h. The stack never holds more than h nodes, so the extra memory is O(h), which means it grows with the height and not with the number of keys.

The lesson uses a plain word, height, for h. In this lesson the height counts nodes. The nine-key example tree has height 4, because the path 8, 3, 6, 4 has four nodes.

The next program measures the peak stack size on two trees. It also builds the full list for comparison. The helper build_from_sorted makes a balanced tree from a sorted array. Its middle key becomes the root, and each half does the same. A chain is a tree where every node has only one child, so the tree is a single long path.

The complete program is at the end of this lesson. Its measured output is:

iterator on balanced 1023 keys: count=1023 peak_stack=10
full list: size=1023 first=0 last=1022
iterator on chain of 1000 keys: peak_stack=1000

Read these numbers one at a time.

• The balanced tree has 1023 keys and 10 levels. The stack peaked at 10 nodes, which is the height. The full list holds all 1023 keys.
• The chain has 1000 keys and height 1000. The stack held 1000 nodes at its peak. The iterator still uses O(h) memory, and here h is large because the tree is a chain.

The chain in this run comes from inserting the keys 999 down to 0. That makes every node the left child of the one above it. If the keys went in ascending order instead, every node would be a right child. The iterator then holds only one node at a time. A quick test showed a peak of 1 for that case. The shape of the tree decides the stack size, not the number of keys alone. The module lesson "Why sorted input makes a BST a chain" explains how such a chain forms.

Time per call

Every node is pushed once and popped once during a full walk. A walk over n keys does about n pushes and n pops, so the total work is O(n). Spread over n calls, the average cost of one next() is O(1). Some calls push many nodes and others push none, so one call can take longer than the average. The average is the number that matters for the whole walk.

Common mistakes

Pushing only the root, and only the right child after each pop. A constructor that pushes just the root skips the left branch. A next() that pushes only the right child skips the left branch of that child too. Both mistakes still print numbers, which makes them easy to miss:

explicit Wrong(const Node* r) { if (r) stack_.push(r); }
 
int next() {
    const Node* n = stack_.top();
    stack_.pop();
    if (n->right) stack_.push(n->right.get());
    return n->key;
}

Run on the nine-key tree, this version prints:

wrong: 8 10 14

The keys 1, 3, 4, 6, 7 and 13 are missing. The walk must go down the whole left branch at the start and after every pop.

Storing key values in the stack instead of node pointers. The stack must hold nodes, not keys. A key does not know its children, so the iterator cannot reach the right branch from it. Here is the mistake:

std::stack<int> s;
s.push(n.key);
std::cout << s.top()->right;

The compiler stops with this message:

mistake4.cpp:9:25: error: base operand of ‘->’ is not a pointer

The message means s.top() is an int, which has no -> operator. The fix is to store const Node* in the stack, as the iterator does.

Calling next() when no key is left. The stack is empty after the last key. Calling top() on an empty stack is an error. With the GNU C++ library that comes with g++ 15.2.0 on this machine, the program stops with this message and exit code 134:

/usr/include/c++/15/bits/stl_stack.h:260: std::stack<_Tp, _Sequence>::reference std::stack<_Tp, _Sequence>::top() [with _Tp = const Node*; _Sequence = std::deque<const Node*, std::allocator<const Node*> >; reference = const Node*&]: Assertion '!this->empty()' failed.

The fix is to check has_next() before each call, as the loops in this lesson do.

Freeing the tree while the iterator still exists. The iterator stores raw addresses of nodes. If the tree is destroyed first, those addresses point to freed memory. The compiler does not catch this, and the result is undefined behaviour, which means the program may crash, print wrong keys, or appear to work. Keep the tree alive for as long as the iterator is used.

The full program

This is the complete program. It is the file used for every output shown in this lesson. It compiles with g++ -std=c++20 -Wall -Wextra -pedantic.

#include <algorithm>
#include <cstddef>
#include <iostream>
#include <memory>
#include <stack>
#include <vector>
 
struct Node {
    int key;
    std::unique_ptr<Node> left;
    std::unique_ptr<Node> right;
    explicit Node(int k) : key(k) {}
};
 
using Tree = std::unique_ptr<Node>;
 
Tree insert(Tree root, int key) {
    if (!root) return std::make_unique<Node>(key);
    if (key < root->key) root->left = insert(std::move(root->left), key);
    else if (key > root->key) root->right = insert(std::move(root->right), key);
    return root;
}
 
Tree build_from_sorted(const std::vector<int>& keys, std::size_t lo, std::size_t hi) {
    if (lo >= hi) return nullptr;
    std::size_t mid = lo + (hi - lo) / 2;
    auto node = std::make_unique<Node>(keys[mid]);
    node->left = build_from_sorted(keys, lo, mid);
    node->right = build_from_sorted(keys, mid + 1, hi);
    return node;
}
 
class BstIterator {
public:
    explicit BstIterator(const Node* root) { push_left(root); }
 
    bool has_next() const { return !stack_.empty(); }
 
    int next() {
        const Node* node = stack_.top();
        stack_.pop();
        push_left(node->right.get());
        return node->key;
    }
 
    std::size_t stack_size() const { return stack_.size(); }
 
private:
    void push_left(const Node* node) {
        for (; node != nullptr; node = node->left.get()) stack_.push(node);
    }
 
    std::stack<const Node*> stack_;
};
 
class DescendingIterator {
public:
    explicit DescendingIterator(const Node* root) { push_right(root); }
 
    bool has_next() const { return !stack_.empty(); }
 
    int next() {
        const Node* node = stack_.top();
        stack_.pop();
        push_right(node->left.get());
        return node->key;
    }
 
private:
    void push_right(const Node* node) {
        for (; node != nullptr; node = node->right.get()) stack_.push(node);
    }
 
    std::stack<const Node*> stack_;
};
 
void collect(const Node* node, std::vector<int>& out) {
    if (!node) return;
    collect(node->left.get(), out);
    out.push_back(node->key);
    collect(node->right.get(), out);
}
 
int main() {
    // Example 1: one tree, keys in sorted order
    Tree root;
    for (int k : {8, 3, 10, 1, 6, 14, 4, 7, 13}) root = insert(std::move(root), k);
    BstIterator it(root.get());
    std::cout << "example 1:";
    while (it.has_next()) std::cout << ' ' << it.next();
    std::cout << '\n';
 
    // Example 2: merge two trees with two iterators
    Tree a, b;
    for (int k : {2, 5, 9}) a = insert(std::move(a), k);
    for (int k : {1, 4, 6, 10}) b = insert(std::move(b), k);
    BstIterator ia(a.get()), ib(b.get());
    bool ha = ia.has_next(), hb = ib.has_next();
    int va = ha ? ia.next() : 0;
    int vb = hb ? ib.next() : 0;
    std::cout << "example 2:";
    while (ha || hb) {
        if (!hb || (ha && va <= vb)) {
            std::cout << ' ' << va;
            ha = ia.has_next();
            if (ha) va = ia.next();
        } else {
            std::cout << ' ' << vb;
            hb = ib.has_next();
            if (hb) vb = ib.next();
        }
    }
    std::cout << '\n';
 
    // Example 3: the same tree, walked from the largest key down
    DescendingIterator di(root.get());
    std::cout << "example 3:";
    while (di.has_next()) std::cout << ' ' << di.next();
    std::cout << '\n';
 
    // Memory: balanced tree of 1023 keys
    std::vector<int> keys(1023);
    for (int i = 0; i < 1023; ++i) keys[i] = i;
    Tree big = build_from_sorted(keys, 0, keys.size());
    BstIterator bi(big.get());
    std::size_t peak = 0;
    std::size_t count = 0;
    while (bi.has_next()) {
        peak = std::max(peak, bi.stack_size());
        bi.next();
        ++count;
    }
    std::cout << "iterator on balanced 1023 keys: count=" << count << " peak_stack=" << peak << '\n';
 
    std::vector<int> all;
    collect(big.get(), all);
    std::cout << "full list: size=" << all.size() << " first=" << all.front() << " last=" << all.back() << '\n';
 
    // Memory: keys inserted in descending order make a chain of left children
    Tree chain;
    for (int i = 999; i >= 0; --i) chain = insert(std::move(chain), i);
    BstIterator ci(chain.get());
    std::size_t cpeak = 0;
    while (ci.has_next()) {
        cpeak = std::max(cpeak, ci.stack_size());
        ci.next();
    }
    std::cout << "iterator on chain of 1000 keys: peak_stack=" << cpeak << '\n';
}

The program prints the following when run:

example 1: 1 3 4 6 7 8 10 13 14
example 2: 1 2 4 5 6 9 10
example 3: 14 13 10 8 7 6 4 3 1
iterator on balanced 1023 keys: count=1023 peak_stack=10
full list: size=1023 first=0 last=1022
iterator on chain of 1000 keys: peak_stack=1000

Practice

Exercise 1. Build the tree from the keys 5, 2, 8, 1, 3 by inserting them in that order. Write the stack after the constructor and after the first call to next(), with the top first.

Show the answer

After the constructor the stack is 1, 2, 5. The constructor walks left from 5 to 2 to 1 and pushes each node. The first call to next() returns 1 and pops it. Node 1 has no right child, so nothing is pushed. The stack is then 2, 5.

Exercise 2. For the tree in exercise 1, what is the largest number of nodes the stack holds at one time during a full walk?

Show the answer

The answer is 3. The stack reaches 3 nodes right after the constructor, when it holds 5, 2 and 1. A full walk over this tree never made it grow past 3.

Exercise 3. Add a method int peek() const to BstIterator. It returns the key that the next call to next() would return, without removing it. Say what the method does when has_next() is false.

Show the answer

int peek() const { return stack_.top()->key; }

The method reads the top node and does not pop it. When the stack is empty, top() fails in the same way as the empty next() call in the mistakes section. The caller must check has_next() first.

Exercise 4. Write a function int kth_smallest(const Node* root, int k) that returns the k-th smallest key, counting from 1. It may use BstIterator. Run it on the nine-key tree with k from 1 to 4.

Show the answer

int kth_smallest(const Node* root, int k) {
    BstIterator it(root);
    int value = 0;
    for (int i = 0; i < k && it.has_next(); ++i) value = it.next();
    return value;
}

On the nine-key tree the results for k = 1, 2, 3 and 4 are 1, 3, 4 and 6. The loop stops after k calls, so the walk does not visit the rest of the tree.

Exercise 5. Build a tree from the keys 1 to 1000, inserted in ascending order. What is the peak stack size of BstIterator during a full walk, and why?

Show the answer

The peak is 1. Ascending insertion makes every node the right child of the one before it, so the tree is a chain going to the right. The constructor pushes only the root, because the root has no left child. After each next() the right child is pushed alone, so the stack never holds more than one node. A quick test with the keys 0 to 999 gave a peak of 1.

Summary

• An inorder walk visits a BST's keys in sorted order, and an iterator produces them one call at a time.
• The stack holds the nodes whose keys are still waiting, and each next() pops one and pushes the left branch of its right child.
• The stack holds only ancestors of the next key, so its size is at most the height of the tree, which is O(h).
• A full walk pushes and pops each node once, so the total work is O(n) and the average cost per call is O(1).
• A full list costs memory for every key and must be built before the first key is returned, while the iterator starts answering right away.