【leetcode热题】二叉搜索树迭代器

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

解法一

先不考虑题目 Note 中要求的空间复杂度和时间复杂度，简单粗暴一些。在构造函数中，对二叉树进行中序遍历，把结果保存到一个队列中，然后 next 方法直接执行出队操作即可。至于 hasNext 方法的话，判断队列是否为空即可。

class BSTIterator {

    Queue<Integer> queue = new LinkedList<>();

    public BSTIterator(TreeNode root) {
        inorderTraversal(root);
    }

    private void inorderTraversal(TreeNode root) {
        if (root == null) {
            return;
        }
        inorderTraversal(root.left);
        queue.offer(root.val);
        inorderTraversal(root.right);
    }

    /** @return the next smallest number */
    public int next() {
        return queue.poll();
    }

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return !queue.isEmpty();
    }
}

时间复杂度的话，构造函数因为遍历了一遍二叉树，所以是 O(n)，对于 next 和 hasNext 方法都是 O(1)。

空间复杂度，用队列保存了所有的节点值，所以是 O(n)，此外中序遍历递归压栈的过程也需要 O(h) 的空间。

解法二

解法一中我们把所有节点都保存了起来，其实没必要一次性保存所有节点，而是需要一个输出一个即可。

所以我们要控制中序遍历的进程，不要让它一次性结束，如果用解法一递归的方法去遍历那就很难控制了，所以自然而然的会想到用栈模拟递归的过程。

下边是 94 题 解法二的代码。

public List<Integer> inorderTraversal(TreeNode root) {
    List<Integer> ans = new ArrayList<>();
    Stack<TreeNode> stack = new Stack<>();
    TreeNode cur = root;
    while (cur != null || !stack.isEmpty()) {
        //节点不为空一直压栈
        while (cur != null) {
            stack.push(cur);
            cur = cur.left; //考虑左子树
        }
        //节点为空，就出栈
        cur = stack.pop();
        //当前值加入
        ans.add(cur.val);
        //考虑右子树
        cur = cur.right;
    }
    return ans;
}

和这道题糅合一起也很简单了，只需要把 stack 和 cur 作为成员变量，然后每次调用 next 就执行一次 while 循环，并且要记录当前值，结束掉本次循环。

class BSTIterator {
    Stack<TreeNode> stack = new Stack<>();
    TreeNode cur = null;

    public BSTIterator(TreeNode root) {
        cur = root;
    }

    /** @return the next smallest number */
    public int next() {
        int res = -1;
        while (cur != null || !stack.isEmpty()) {
            // 节点不为空一直压栈
            while (cur != null) {
                stack.push(cur);
                cur = cur.left; // 考虑左子树
            }
            // 节点为空，就出栈
            cur = stack.pop();
            res = cur.val;
            // 考虑右子树
            cur = cur.right;
            break;
        }

        return res;
    }

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return cur != null || !stack.isEmpty();
    }
}

时间复杂度的话，对于 next 方法，大多数时候是 O(1)，但最坏情况因为最里边的 while 循环，其实有可能达到 O(n)。但如果算均摊时间复杂度的话，其实还是 O(1)，因为每个节点最多也就经过两次就出栈了。

空间复杂度，这里只需要消耗栈的空间，也就是 O(h)。