BST Iterator with Stack Traversal

BST Iterator with Stack Traversal

My first approach was pretty brute force. I built a vector to store all node values in ascending order.

class BSTIterator {
public:
    BSTIterator(TreeNode* root) {
        if (root==NULL) return;
        cur = 0;
        build_tree(root);
    }
    
    void build_tree(TreeNode* root) {
        if (root->left) build_tree(root->left);
        vals.push_back(root->val);
        if (root->right) build_tree(root->right);
    }
    
    /** @return the next smallest number */
    int next() {
        return vals[cur++];
    }
    
    /** @return whether we have a next smallest number */
    bool hasNext() {
        if (vals.size()==0) return false;
        return cur < vals.size();
    }
    
    int cur;
    vector<int> vals;
};

The better approach uses a stack to traverse left subtrees continuously. This gives us O(1) amortized time complexity.

class BSTIterator {
public:
    BSTIterator(TreeNode* root) {
        if (root==NULL) return;
        build_tree(root);
    }
    
    void build_tree(TreeNode* root) {
        if (root==NULL) return;
        node.push(root);
        while(root->left) {
            root=root->left;
            node.push(root);
        }
    }
    
    /** @return the next smallest number */
    int next() {
        TreeNode* tmp = node.top();
        node.pop();
        build_tree(tmp->right);
        return tmp->val;
    }
    
    /** @return whether we have a next smallest number */
    bool hasNext() {
        return !node.empty();
    }

    stack<TreeNode*> node;
};

#stack #BST