Binary Search Tree Iterator

Question (LC.173)

Implement an iterator over a binary search tree.

It will be initialized with the root node of a BST.

Example

// root is the root node of BST [4,2,10,1,3,6,11]
BSTIterator iter = new BSTIterator(root);
while (iter.hasNext()) print iter.next() + ",";
// output: 1,2,3,4,6,10,11

Approach

Iterative inorder traversal

Code

public class BSTIterator {
    private Stack<TreeNode> stack;
    private TreeNode current;

    /** Initialize the stack */
    public BSTIterator(TreeNode root) {
        stack = new Stack<TreeNode>();
        current = root;
    }

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        // we are not at the bottom right of the BST
        return (!stack.isEmpty() || current != null);
    }

    /** @return the next smallest number */
    public int next() {
        int smallest;
        // left child first
        while (current != null) {
            stack.push(current);
            current = current.left;
        } 
        // backtrack to parent
        // log inorder
        current = stack.pop();
        smallest = current.val;
        // right child last
        current = current.right;
        return smallest;
    }
}

To some degree, each next() is kind of like a recursive call, given a new root and perform the in-order traversal from there.

class BSTIterator:

    def __init__(self, root: TreeNode):
        self.stack = deque()
        self.cur = root
        

    def next(self) -> int:
        # keep dumping
        while self.cur is not None:
            self.stack.append(self.cur)
            self.cur = self.cur.left
        # pop
        self.cur = self.stack.pop()
        next_val = self.cur.val
        self.cur = self.cur.right
        return next_val

    def hasNext(self) -> bool:
        return len(self.stack) > 0 or self.cur is not None

Reading

Traversing binary trees simply and cheaply by Joseph M. Morris