173 Binary Search Tree Iterator

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.

Calling next() will return the next smallest number in the BST.

Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.

The Idea: Follow an iterative traversal for in order traversal. An in order traversal follows LVR, (left, value, right) recursively. That is, the next smallest element is always going to be absolute left of the tree and then followed right once right once before it attempts to go left again. Following this idea, first initialize the stack with all the left elements. It will then follow that the top of the stack will have the from recent, or leftest element. This is the element we want to return, but our next round needs to be accounted for. Go the right once, and then append all the left elements of this right node (including the right node itself). Now the right node is in the stack, and so are the next set of left elements.

Complexity: O(1) time for hasNext() and next() and O(|height|) space maintained by the stack.

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
public:
    BSTIterator(TreeNode *root) {
        fill_stack(root);
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return !s.empty();
    }

    /** @return the next smallest number */
    int next() {
        TreeNode *top = s.top(); s.pop();
        fill_stack(top->right);
        return top->val;
    }

private:
    stack<TreeNode*> s;

    void fill_stack(TreeNode *root) {
        while (root) {
            s.push(root);
            root = root->left;
        }
    }
};

/**
 * Your BSTIterator will be called like this:
 * BSTIterator i = BSTIterator(root);
 * while (i.hasNext()) cout << i.next();
 */

Last updated