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.

Solution: Using stack

We need something to keep track of the path from root to the current node, so we can move to the parent when needed. Do note that storing the path from root to the current node only requires memory equivalent to the length of the path which is the depth of the tree. Also, we can track the path using stack.

class BSTIterator {
public:
    BSTIterator(TreeNode *root) {
        pushAll(root);
    }

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

    // return the next smallest number
    int next() {
        TreeNode *temp = myStack.top();
        myStack.pop();
        pushAll(temp->right);
        return temp->val;
    }

private:
    stack<TreeNode *> myStack;

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

Solution: Using inorder successor

Time Complexity: O(h) for hasNext() and next()

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

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return inorderSuccessor(root, cur) != NULL;
    }

    /** @return the next smallest number */
    int next() {
        cur = inorderSuccessor(root, cur);
        return cur->val;
    }

private:
    TreeNode *root;
    TreeNode *cur;

    TreeNode *inorderSuccessor(TreeNode *root, TreeNode *p) {
        if (!root) return NULL;

        if (!p) {
            TreeNode *node = root;
            while (node->left) node = node->left;
            return node;
        }

        TreeNode *node = root, *succ = NULL;
        while (node) {
            if (p->val < node->val) {
                succ = node;
                node = node->left;
            } else {
                node = node->right;
            }
        }
        return succ;
    }
};