【Leetcode】173.二叉搜索树迭代器

173.二叉搜索树迭代器

题目

实现一个二叉搜索树迭代器。你将使用二叉搜索树的根节点初始化迭代器。

调用 next() 将返回二叉搜索树中的下一个最小的数。 

示例：

    7
  /   \
 3     15
      /   \
     9    20 
BSTIterator iterator = new BSTIterator(root);
iterator.next();    // 返回 3
iterator.next();    // 返回 7
iterator.hasNext(); // 返回 true
iterator.next();    // 返回 9
iterator.hasNext(); // 返回 true
iterator.next();    // 返回 15
iterator.hasNext(); // 返回 true
iterator.next();    // 返回 20
iterator.hasNext(); // 返回 false

要求：

  • next() 和 hasNext() 操作的时间复杂度是 O(1)，并使用 O(h) 内存，其中 h 是树的高度。
  • 你可以假设 next() 调用总是有效的，也就是说，当调用 next() 时，BST 中至少存在一个下一个最小的数。

方法一（用队列）

我们知道对于二叉搜索树来说，其中序遍历的结果是一个升序序列。因此，我们先中序遍历二叉搜索树，将树中元素的升序排列存放在一个队列中。每当调用 next() 时，我们就从队列中弹出一个元素（升序保证它是最小的），每当调用hasNext() 时，我们只需检查队列中还有没有剩余元素即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class BSTIterator {
    private Queue<Integer> queue;
    public BSTIterator(TreeNode root) {
        queue = new LinkedList<>();
        dfs(root);
    }
    private void dfs(TreeNode root){
        if(root == null)
            return;
        dfs(root.left);
        queue.add(root.val);
        dfs(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();
    }
}

对于这种方法来说，next和hasNext函数的时间复杂度都为O(1),符合题目要求。但是需要用O(n)辅助空间的队列来存储元素，不符合题目中空间复杂度为O(h)的要求（h为树高度）

方法二（用栈）

对于二叉搜索树来说，其最左的那个元素就是树中最小的元素。因此，我们先将根节点root及其左子树的所有节点依次入栈，这样栈顶就是树中最左的元素，也即树中最小的元素。

每当调用next()，返回栈顶元素（树中最小元素），并将其右子节点及右子节点左子树的所有节点入栈，以便之后的next()操作。

每当调用hasNext()时，我们只需检查队列中还有没有剩余元素即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class BSTIterator {
    private Stack<TreeNode> stack;
    public BSTIterator(TreeNode root) {
        stack = new Stack<>();
        leftInorder(root);
    }
    //将root及其子树的所有节点入栈
    private void leftInorder(TreeNode root){
        while(root != null){
            stack.push(root);
            root = root.left;
        }
    }
    
    /** @return the next smallest number */
    public int next() {
        TreeNode mostleftNode = stack.pop();
        if(mostleftNode.right != null)
            leftInorder(mostleftNode.right);
        return mostleftNode.val;
    }
    
    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return stack.size() > 0;
    }
}

时间复杂度：

  • next(): 平均下来为O(1)
  • hasNext(): O(1)

空间复杂度：O(h) h为树的高度