Binary Search Tree Iterator

https://leetcode.com/problems/binary-search-tree-iterator/description/

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.

Thoughts

实现BST的迭代器。看到BST就知道要利用它中序遍历是从小到大有序的性质。巧妙利用栈和push_left: next()时让栈抛出顶点t作为返回值,对它的右结点调用push_left。当右子树全部遍历完,栈会自动退到t的parent。以此类推。

Code

class BSTIterator {
public:
    stack<TreeNode*> s;
    int c = 0;

    void push_left(TreeNode *cur) {
        while (cur != nullptr) {
            s.push(cur);
            cur = cur->left;
        }
    }

    BSTIterator(TreeNode* root) {
        push_left(root);
    }
    
    /** @return the next smallest number */
    int next() {
        const auto t = s.top();
        s.pop();
        push_left(t->right);
        return t->val;
    }
    
    /** @return whether we have a next smallest number */
    bool hasNext() {
        return !s.empty();
    }
};

Analysis

时间复杂度是avg O(1),因为每个结点共进出一次栈。

Last updated