Binary Search Tree Iterator

How it works

This is a slight variation of the previous approach. Instead of initializing the stack in the constructor, we defer the leftward traversal to the next() method. We keep a pointer to the current node and only push nodes onto the stack when next() is called. This makes the constructor O(1) but the logic is essentially the same.

Steps

1. In the constructor, store the root as the current node and initialize an empty stack.
2. For next(), push all nodes from the current node down to its leftmost descendant onto the stack.
3. Pop the top node from the stack, set the current node to its right child.
4. Return the popped node's value.
5. For hasNext(), return true if either the current node is not null or the stack is not empty.

Complexity

• Time complexity: O(1) in average for each function call.
• Space complexity: O(h)

Where n is the number of nodes and h is the height of the given tree.

Things to watch for

Using Preorder or Postorder Instead of Inorder

BST iterator must return elements in ascending order, which requires inorder traversal (left, root, right). Using preorder or postorder traversal produces elements in the wrong order.

Forgetting to Traverse the Right Subtree After Popping

When using the stack-based approach, after popping a node and returning its value, you must push the leftmost path of its right subtree onto the stack. Forgetting this step causes right subtrees to be skipped entirely.

Incorrect hasNext() Implementation

The hasNext() method must return True if there are more elements to iterate. For the stack-based approach, this means checking if either the stack is non-empty OR the current pointer is not null (depending on the implementation variant).

from typing import Optional

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.cur = root
        self.stack = []

    def next(self) -> int:
        while self.cur:
            self.stack.append(self.cur)
            self.cur = self.cur.left

        node = self.stack.pop()
        self.cur = node.right
        return node.val

    def hasNext(self) -> bool:
        return bool(self.cur) or bool(self.stack)

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class BSTIterator {
private:
    TreeNode* cur;
    stack<TreeNode*> stack;

public:
    BSTIterator(TreeNode* root) {
        cur = root;
    }

    int next() {
        while (cur) {
            stack.push(cur);
            cur = cur->left;
        }

        TreeNode* node = stack.top();
        stack.pop();
        cur = node->right;
        return node->val;
    }

    bool hasNext() {
        return cur || !stack.empty();
    }
};