173. Binary Search Tree Iterator

173. Binary Search Tree Iterator

刚开始的思路比较暴力,直接构建一个 vector 来从小到大存节点的值

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;
};

更好的思路是使用栈来不断地遍历左子树,均摊计算使得平均时间复杂度为 O(1)

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