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.

Credits:
Special thanks to @ts for adding this problem and creating all test cases.

A cleaner solution, 7/10/2015 update

Use a stack to track the path in the tree.

Keep this invariant in mind, s.top() is always the next node to visit.

if s is empty, hasNext return false.

when an element is popped, check if its right child is null. if its right child exists, push it in to the stack, and push its left child and move to left child recursively.

It’s an iterated implementation of in order traversal.

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
public:
    stack<TreeNode *> s;
    BSTIterator(TreeNode *root) {
        //invariant s.top() is the next node, if s.empty() hasNext return false;
        while(root){
            s.push(root);
            root = root->left;
        }
        
    }

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

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

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

Use a stack to store the current path in the tree.

Pop an element in the stack when it is useless. (The right child of the node is visited.)

Be aware of NULL input and start state.

//force solution: convert the BST to a sequence. O(n) total in time, and O(n) in space
/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
public:
    stack<TreeNode *> path;
    bool isStart;
    bool isEnd;
    BSTIterator(TreeNode *root) {
        if(root == NULL){
            isEnd = true;
            isStart = true;
        }
        else{
            isEnd = false;
            path.push(root);
            isStart = false;
        }
        
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        if(isEnd == true){
            return false;
        }
        if(isStart == false){
            return true;
        }
        else if(path.top()->right != NULL || path.size() > 1){
            return true;
        }
        return false;
    }

    /** @return the next smallest number */
    int next() {
        if(path.empty()) return 0;
        TreeNode * p = path.top();
        if(isStart == false){
            isStart = true;
            while(p->left != NULL){
                p = p->left;
                path.push(p);
            }
            return p->val;
        }
        else{
            if(p->right != NULL){
                //pop the top, cause it will never be used
                path.pop();
                p = p->right;
                path.push(p);
                while(p->left != NULL){
                    p = p->left;
                    path.push(p);
                }
                return p->val;
            }
            else{
                path.pop();
                return path.top()->val;
            }
        }
    }
};

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