LeetCode 173. 二叉搜索树迭代器

2023-03-21 HGNULB 约 5 分钟

文章信息 10 个分类 6 个标签

分类

LeetCode CodeTop 中等未解决 5-10 字节跳动微软快手百度面试经典 150 题

题目描述

✅ 173. 二叉搜索树迭代器

题意分析

这是一道设计题：给定一棵二叉搜索树的根，要求实现一个对象，对外提供「取下一个值」和「是否还有下一个值」两个操作，且取出的值必须按从小到大的顺序排列。

题目明确要求 next 和 hasNext 的均摊时间为 $O(1)$，额外空间为 $O(h)$，其中 $h$ 是树高。这两个指标合起来是整道题的核心信号：$O(h)$ 的空间上限直接封死了「构造时把整棵树的值抖成一个数组」的做法，因为那需要 $O(n)$ 的存储；而均摊 $O(1)$ 的要求又说明允许某一次调用偶尔慢一些，只要总量摊得平即可。换句话说，题目要的是一个「按需展开、边走边算」的遍历过程。

调用方保证只在 hasNext 为真时才调用 next，所以不必为「取空」设计异常路径。树中节点值互不相同，也不必处理相等值的先后顺序。边界情形是根为空的树，此时 hasNext 从一开始就应当返回假。

解法：受控中序遍历

核心思路

二叉搜索树的中序遍历天然按升序访问节点。若构造时直接展开整棵树，next() 虽然是 $O(1)$，却需要 $O(n)$ 空间，不满足题目要求。更好的做法是把中序遍历暂停在“下一个节点”之前，按需继续。

栈中只保存从某个待访问子树根到其最左节点的路径，并维护不变量：

栈顶是当前尚未访问节点中的最小值；栈内节点的左子树已经处理完，自身和右子树尚未处理。

构造时压入根的整条左链。next() 弹出栈顶后，该节点的下一个候选来自它的右子树，因此再把右子树的左链压栈。不变量随之恢复。hasNext() 只需判断栈是否为空，且不能推进遍历状态。

每个节点一生只入栈、出栈各一次，所以某次 next() 虽可能压入一条长左链，但连续调用 $n$ 次的总成本是 $O(n)$，均摊为 $O(1)$。

解题步骤

构造函数调用 pushLeft(root)，把根到最左节点的路径压栈。

next() 弹出栈顶；由不变量，它就是当前最小的未访问节点。

对弹出节点的右孩子调用 pushLeft，补入右子树中接下来应访问的路径。

返回弹出节点的值。

hasNext() 返回栈是否非空。

例如 [7,3,15,null,null,9,20]：初始栈顶为 3；弹出 3 后栈顶为 7；弹出 7 时压入右子树左链 15 -> 9，新栈顶变为 9。输出依次为 3, 7, 9, 15, 20。

代码实现

import java.util.ArrayDeque;
import java.util.Deque;

class BSTIterator {
    private final Deque<TreeNode> stack = new ArrayDeque<>();

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

    public int next() {
        TreeNode node = stack.pop();
        pushLeft(node.right);
        return node.val;
    }

    public boolean hasNext() {
        return !stack.isEmpty();
    }

    private void pushLeft(TreeNode node) {
        while (node != null) {
            stack.push(node);
            node = node.left;
        }
    }
}

type BSTIterator struct {
    stack []*TreeNode
}

func Constructor(root *TreeNode) BSTIterator {
    iterator := BSTIterator{}
    iterator.pushLeft(root)
    return iterator
}

func (this *BSTIterator) Next() int {
    last := len(this.stack) - 1
    node := this.stack[last]
    this.stack = this.stack[:last]
    this.pushLeft(node.Right)
    return node.Val
}

func (this *BSTIterator) HasNext() bool {
    return len(this.stack) > 0
}

func (this *BSTIterator) pushLeft(node *TreeNode) {
    for node != nil {
        this.stack = append(this.stack, node)
        node = node.Left
    }
}

复杂度分析

时间复杂度：构造为 $O(h)$；hasNext() 为 $O(1)$；next() 单次最坏为 $O(h)$，但每个节点只入栈、出栈一次，因此均摊为 $O(1)$。

空间复杂度：$O(h)$，其中 $h$ 是树高。栈始终只保存一条搜索路径。

关键点总结

显式栈保存的是递归中序遍历尚未完成的调用现场。

“只压左链”让栈顶始终对应下一个最小值。

均摊分析要按完整遍历记账：总入栈次数和总出栈次数都不超过 $n$。

hasNext() 必须只读；重复调用不应改变下一次 next() 的结果。

易错点总结

构造时只压根节点：第一次返回的未必是最小值。

弹出节点后忘记处理右子树：右侧节点会被全部跳过。

对右孩子只压一个节点：会漏掉右子树更小的左链节点。

让 hasNext() 顺便推进游标：连续调用会跳过元素。

Go 使用值接收者实现 Next()：栈的截断不会保存到迭代器状态中。

相似题目

题目	难度	考察点
LCR 055. 二叉搜索树迭代器	中等	同题换编号，可直接复用同一份实现
94. 二叉树的中序遍历	简单	一次性输出整个序列，无需拆成可暂停的接口
285. 二叉搜索树中的中序后继	中等	只求单个节点的后继，可利用搜索性质直接下降
341. 扁平化嵌套列表迭代器	中等	同样用显式栈做惰性展开，但结构是嵌套列表
230. 二叉搜索树中第 K 小的元素	中等	只需推进 k 步即可停止，是本迭代器的典型用法