《剑指 Offer》专项突破版 - 面试题 55 : 二叉搜索树迭代器（C++ 实现的三种方法）

目录

前言

方法一

方法二

方法三

---

 

---

前言

题目链接：LCR 055. 二叉搜索树迭代器 - 力扣（LeetCode）

题目：

请实现二叉搜索树的迭代器 BSTIterator，它主要有如下 3 个函数。

• 构造函数：输入二叉搜索树的根节点初始化该迭代器。

• 函数 next：返回二叉搜索中下一个最小的节点的值。

• 函数 hasNext：返回二叉搜索树是否还有下一个节点。

分析：

例如，输入下图中的二叉搜索树初始化 BSTIterator，第 1 次调用函数 next 将返回最小的节点的值 1，此时调用函数 hasNext 返回 true。第 2 次调用函数 next 将返回下一个节点的值 2，此时再调用函数 hasNext 将返回 true。第 3 次调用函数 next 将返回下一个节点的值 3，此时再调用函数 hasNext 将返回 false。

---

方法一

如果允许修改输入的二叉搜索树，则可以在初始化迭代器时将它展平成除叶节点之外其他节点只有一个右子节点。这时二叉搜索树看起来像一个链表。然后用一个指针 P 指向展平后的二叉搜索树的根节点，如果指针 P 指向的节点不为 nullptr，那么函数 hasNext 将返回 true。每次函数 next 被调用时都返回指针 P 指向节点的值，并将指针 P 朝着指向右子节点的指针向前移动一次。面试题 52 详细讨论了如何将一棵二叉搜索树展平。

class BSTIterator {
public:
    BSTIterator(TreeNode* root) {
        stack<TreeNode*> st;
        TreeNode* prev = nullptr;
        TreeNode* cur = root;
        while (cur || !st.empty())
        {
            while (cur)
            {
                st.push(cur);
                cur = cur->left;
            }
​
            cur = st.top();
            st.pop();
            cur->left = nullptr;
            if (prev != nullptr)
            {
                prev->right = cur;
            }
            else
            {
                p = cur;
            }
            prev = cur;
            cur = cur->right;
        }
    }
    
    int next() {
        int result = p->val;
        p = p->right;
        return result;
    }
    
    bool hasNext() {
        return p != nullptr;
    }
private:
    TreeNode* p;
};

函数 next 和 hasNext 的时间复杂度以及空间复杂度都是 O(1)。

---

方法二

如果不允许修改输入的二叉搜索树，那么可以在迭代器初始化时另外创建一个数组保存二叉树所有节点的值。

class BSTIterator {
public:
    BSTIterator(TreeNode* root) : i(0) {
        stack<TreeNode*> st;
        TreeNode* cur = root;
        while (cur || !st.empty())
        {
            while (cur)
            {
                st.push(cur);
                cur = cur->left;
            }
​
            cur = st.top();
            st.pop();
            v.push_back(cur->val);
            cur = cur->right;
        }
    }
    
    int next() {
        return v[i++];
    }
    
    bool hasNext() {
        return i < v.size();
    }
private:
    vector<int> v;
    int i;
};

函数 next 和 hasNext 的时间复杂度都是 O(1)，但空间复杂度都是 O(n)。

---

方法三

可以在不修改输入的二叉搜索树的同时优化空间效率。如果对二叉搜索的中序遍历的迭代代码足够熟悉，我们就会注意到中序遍历的迭代代码中有一个 while 循环，循环的条件为 true 时循环体每次执行一次就遍历二叉树的一个节点。当 while 循环的条件为 false 时，二叉树中所有节点都已遍历完。因此，中序遍历的迭代代码中的 while 循环可以看成迭代器 hasNext 的判断条件，而 while 循环体内执行的操作就是函数 next 执行的操作。

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

函数 next 和 hasNext 的时间复杂度都是 O(1)，空间复杂度都是 O(h)，其中 h 为二叉搜索树的深度。