【LeetCode每日一题】173. 二叉搜索树迭代器

【LeetCode每日一题】173. 二叉搜索树迭代器

1.题目

实现一个二叉搜索树迭代器类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

 

提示：

树中节点的数目在范围 [1, 105] 内 0 <= Node.val <= 106 最多调用 105 次 hasNext 和 next 操作

2.实现

本题考察二叉树中序遍历非递归解法，可以采用预处理，以递归方法为例，将二叉树进行中序遍历并将遍历的结点存储下来，调用函数直接取。假设结点有n个，时间与空间复杂度均为O(n)。下面代码实现采用了两种方式进行预处理。

class BSTIterator {
public:
    vector<int> v;
    int i;
    BSTIterator(TreeNode* root) {
        i = 0;
        inorderNR(root); // 预处理 可以选择递归与非递归
    }
    // 递归
    void inorder(TreeNode* root) {
        if (!root) return;
        inorder(root->left);
        v.push_back(root->val);
        inorder(root->right);
    }
    // 非递归
    void inorderNR(TreeNode* root) {
        TreeNode* cur = root;
        stack<TreeNode*> st;
        while (cur || !st.empty()) {
            while(cur) {
                st.push(cur);  // 栈底->栈顶顺序为 根->左
                cur = cur->left;
            }
            // 栈顶存储的是最左侧结点
            if(!st.empty()) {
                auto r = st.top(); st.pop();
                v.push_back(r->val);
                cur = r->right; 
            }
        }
    }
    
    int next() {
        return v[i++];
    }
    
    bool hasNext() {
        return i < v.size();
    }
};

在进阶中我们可以看到要求是：next() 和 hasNext() 操作均摊时间复杂度为 O(1) ，并使用 O(h) 内存。其中 h 是树的高度。

我们可以采用非递归方式，将其融入到next函数中，不做预先处理，实现如下：

class BSTIterator {
public:
    TreeNode* cur; // 左子数的最右结点
    stack<TreeNode*> st;
    BSTIterator(TreeNode* root) {
        cur = root;
    }
    
    int next() {
        while(cur) {
            st.push(cur);  // 栈底->栈顶顺序为 根->左
            cur = cur->left;
        }
        auto r = st.top(); st.pop();
        cur = r->right; 
        return r->val;
    }
    
    bool hasNext() {
        return cur || !st.empty();
    }
};

上面需要采用cur保存左子树的最右结点，而另外一种栈写法是直接模拟中序遍历，由于顺序是左->根->右，那么栈的顺序就是右->根->左。

class BSTIterator {
public:
    struct Command {
        string flag;
        TreeNode* node;
    };
    stack<Command> st;
    BSTIterator(TreeNode* root) {
        st.push({"visit",root});
    }
    
    int next() {
        int res;
        while (true) {
            auto x = st.top(); st.pop();
            if (x.flag == "print") {
                res = x.node->val;
                break;
            } else {
                if (x.node->right) {
                    st.push({"visit",x.node->right});
                }
                st.push({"print",x.node});
                if (x.node->left) {
                    st.push({"visit",x.node->left});
                }
            }
        }
        return res;
    }
    
    bool hasNext() {
        return !st.empty();
    }
};