173. 二叉搜索树迭代器

实现一个二叉搜索树迭代器类BSTIterator，表示一个按中序遍历二叉搜索树（BST）的迭代器： BSTIterator(TreeNode root)初始化 BSTIterator类的一个对象。BST的根节点 root会作为构造函数的一部分给出。指针应初始化为一个不存在于 BST中的数字，且该数字小于 BST 中的任何元素。 boolean hasNext()如果向指针右侧遍历存在数字，则返回 true；否则返回 false。 int next()将指针向右移动，然后返回指针处的数字。 注意，指针初始化为一个不存在于 BST中的数字，所以对next()的首次调用将返回 BST中的最小元素。

你可以假设 next() 调用总是有效的，也就是说，当调用next() 时，BST的中序遍历中至少存在一个下一个数字。

示例：

输入

["BSTIterator", "next", "next", "hasNext", "next", "hasNext", "next", "hasNext", "next", "hasNext"]

[[[7, 3, 15, null, null, 9, 20]], [], [], [], [], [], [], [], [], []]

输出

[null, 3, 7, true, 9, true, 15, true, 20, false]

解释

BSTIterator bSTIterator = new BSTIterator([7, 3, 15, null, null, 9, 20]);

bSTIterator.next(); // 返回 3

bSTIterator.next(); // 返回 7

bSTIterator.hasNext(); // 返回 True

bSTIterator.next(); // 返回 9

bSTIterator.hasNext(); // 返回 True

bSTIterator.next(); // 返回 15

bSTIterator.hasNext(); // 返回 True

bSTIterator.next(); // 返回 20

bSTIterator.hasNext(); // 返回 False

思路：

方法一：扁平化

我们可以直接对二叉搜索树做一次完全的递归遍历，获取中序遍历的全部结果并保存在数组中。随后，我们利用得到的数组本身来实现迭代器。

时间复杂度：O(n)

空间复杂度：O(n)

class BSTIterator {
private:
    vector<int> arr;
    int idx;
    void inorder(TreeNode* root, vector<int>& res) {
        if(!root){
            return;
        }
        inorder(root->left, res);
        res.push_back(root->val);
        inorder(root->right, res);
    }
    vector<int> inorderTraversal(TreeNode* root){
        vector<int> res;
        inorder(root, res);
        return res;
    }
public:
    BSTIterator(TreeNode* root): idx(0), arr(inorderTraversal(root)) {}
    
    int next() {
        return arr[idx++];
    }
    
    bool hasNext() {
        return (idx < arr.size());
    }
};

方法二：迭代

利用栈，通过迭代的方式对二叉树做中序遍历。此时，我们无需预先计算出中序遍历的全部结果，只需要实时维护当前栈的情况即可。

时间复杂度：O(n)

空间复杂度：O(n)，取决于栈的深度，最坏情况为一条链。

class BSTIterator1 {
public:
    TreeNode* cur;
    stack<TreeNode*> stk;
    BSTIterator(TreeNode* root): cur(root) {}
    
    int next() {
        while(cur != nullptr){
            stk.push(cur);
            cur = cur->left;
        }
        cur = stk.top();
        stk.pop();
        int ret = cur->val;
        cur = cur->right;
        return ret;
    }
    
    bool hasNext() {
        return cur != nullptr || !stk.empty();
    }
};