LeetCode LCR 055. 二叉搜索树迭代器

约 8 分钟
文章信息 5 个分类 7 个标签
分类
LeetCode 中等 119 经典题变种挑战 剑指 Offer 剑指 Offer（专项突击版）
标签
树 二叉树 LCR 栈 设计 二叉搜索树 迭代器
目录

题目描述

✅ LCR 055. 二叉搜索树迭代器

题意分析

题目目标：为二叉搜索树实现一个迭代器，next 按升序返回下一个节点值，hasNext 判断是否还有元素，构造函数接收树根。
核心约束：升序等价于中序，所以本质是把一次中序遍历"拆开"成可以随时暂停、随时继续的形式。题目还明确要求 next 和 hasNext 均摊 $O(1)$、空间 $O(h)$，这条要求直接否决了"构造时把所有值摊平进数组"的写法——那样空间是 $O(n)$。
边界处理：树至少有一个节点，但调用方可能在取完所有元素后继续调用 hasNext（必须返回假）；最小值可能藏在很深的左链末端；某个节点没有右子树时，下一个元素要回到祖先。
实现取舍：递归的中序遍历没法在中途"停下来"，因为进度保存在调用栈里而调用栈不受我们控制。要实现可暂停，就必须把这个调用栈显式地拿到手上自己管理。

解法：树形遍历

核心思路

最省事的写法是构造时做一次完整中序遍历，把所有值存进数组，next 返回下标处的元素并自增。这样两个操作确实是 $O(1)$，但空间是 $O(n)$，直接违反题目对空间的要求；而且如果调用方只取前几个元素，前期就白白付出了遍历整棵树的代价。
观察递归中序的执行过程：它总是先沿左孩子一路下探到最左端，再逐个回溯。这个"下探路径"就是递归调用栈的内容，也正是遍历进度的全部信息。既然递归的栈不可控，那就把它换成一个我们自己持有的显式栈——状态存在对象里，方法返回后依然保留，暂停与继续就都成立了。
由此确定不变量：任意时刻栈中自底向上保存的，是"下一个待返回节点"及其所有尚未被返回的祖先，且栈顶恰好就是当前最小的未返回节点。构造时从根一路向左压栈即可建立这个不变量。
维护它靠 next 里的两步：弹出栈顶 cur 作为本次答案（它的左子树已经全部返回完毕），然后把 cur 的右子树的左链整条压栈——因为中序里紧随 cur 之后的正是它右子树中的最小节点；若右子树为空则什么都不压，栈顶自然回退到祖先，恰好就是下一个应返回的元素。
hasNext 只需判断栈是否为空：不变量保证栈非空当且仅当还有未返回的节点。

解题步骤

  • 构造函数从根出发沿左孩子一路压栈，直到走到空。为什么这样就能让栈顶是最小值：二叉搜索树中最左端的节点即全局最小，压栈路径上的节点则是它的全部祖先，正好符合不变量。
  • hasNext 返回"栈非空"。为什么不需要额外计数：不变量已经把"是否还有元素"这件事编码进了栈的空与非空，多余的计数器反而是新的错误来源。
  • next 先弹出栈顶节点 cur。为什么栈顶就是答案：它的左子树在此前的压栈过程中已经全部处理完毕，按中序它就是当前最小的未返回节点。
  • 弹出后把 cur.right 的左链整条压入栈。为什么是右子树的左链：中序在访问完一个节点后转向它的右子树，而右子树里最先被访问的是其最左端节点，把这条链压进去正好让栈顶指向它。
  • 为什么右子树为空时什么都不做就是正确的：此时中序的下一个节点是"最近的、把 cur 放在其左子树里的祖先"，而这个祖先恰好就是弹出 cur 之后新的栈顶，不需要任何额外操作。
  • 最后返回 cur.val。为什么可以在压栈之后再返回：cur 已经从栈里取出并保存在局部变量中，后续压栈不会影响它。
  • 以 具体用例：树 [7, 3, 15, null, null, 9, 20] 走一遍。构造时从 7 向左压入 7、3，栈自底向上是 [7, 3]。第一次 next：弹出 3，它没有右子树故不压栈，栈变 [7]，返回 3。第二次 next：弹出 7，把右子树 15 的左链 15、9 压入，栈变 [15, 9]，返回 7。此时 hasNext 为真。第三次 next：弹出 9，无右子树，栈变 [15]，返回 9。第四次 next：弹出 15，把右子树 20 的左链压入，栈变 [20]，返回 15。第五次 next：弹出 20，无右子树，栈空，返回 20。此后 hasNext 返回假。输出序列 3, 7, 9, 15, 20 正是中序升序，且栈的最大高度只有 2，从未达到节点总数。

代码实现

// 核心实现：树形遍历，维护必要状态并避免重复处理。
class BSTIterator {
    private Deque<TreeNode> stack = new ArrayDeque<>();

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

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

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

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

// 核心实现：树形遍历，维护必要状态并避免重复处理。
type BSTIterator struct {
    stack []*TreeNode
}

func Constructor(root *TreeNode) BSTIterator {
    var stack []*TreeNode
    for ; root != nil; root = root.Left {
        stack = append(stack, root)
    }
    return BSTIterator{
        stack: stack,
    }
}

func (this *BSTIterator) Next() int {
    cur := this.stack[len(this.stack)-1]
    this.stack = this.stack[:len(this.stack)-1]
    for node := cur.Right; node != nil; node = node.Left {
        this.stack = append(this.stack, node)
    }
    return cur.Val
}

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

复杂度分析

  • 时间复杂度：hasNext 为 $O(1)$，next 均摊 $O(1)$（单次最坏 $O(h)$）。凭什么：每个节点在整个生命周期里只会被压栈一次、弹栈一次，n 次 next 的总压栈次数不超过 n，平均到每次调用就是常数。
  • 空间复杂度：$O(h)$，h 为树高。凭什么：栈里存的始终是一条从根往下的路径上的部分节点，深度不超过树高；相比预先摊平的 $O(n)$，这正是题目要求的那一档。

关键点总结

  • 把递归改成显式栈，是实现"可暂停遍历"的通用手段：递归的进度存在不可控的调用栈里，而迭代器要求进度必须存在对象里，这条动机比代码本身更值得说清楚。
  • 中序迭代器的不变量一句话就能讲明白——"栈顶是下一个待返回节点，栈内其余是它尚未返回的祖先"；只要每个方法执行后这句话仍然成立，实现就是对的。
  • 均摊分析是这类结构的标配论据：单次 next 可能压入一整条左链，但每个节点一生只被压一次，总量线性。
  • 右子树为空时无需任何补救，栈自动回退到正确的祖先——这一点是显式栈写法比"记录父指针"写法优雅的地方。
  • 面试视角：面试官问这题就是冲着"$O(h)$ 空间 + 均摊 $O(1)$"来的。开口先否定预先摊平数组的方案并说明理由，再给出受控栈实现和均摊证明；如果被追问"能否支持 prev"，可以提到需要同时维护前驱方向的栈或改用 Morris 变体，思路展示到这一层就足够了。

易错点总结

  • 错误写法：构造时把整棵树中序摊平进数组 → 树有 $10^5$ 个节点而调用方只取前 3 个时，仍付出 $O(n)$ 的时间与空间，直接违反题目对 $O(h)$ 空间的要求。
  • 错误写法：next 中压栈的是 cur.left 而不是 cur.right 的左链 → 树 [7, 3, 15] 中弹出 3 后又把它的左子树压回，同一节点被重复返回甚至死循环。
  • 错误写法：next 里压的是整条右链（一路向右）而不是右子树的左链 → 树 [7, 3, 15, null, null, 9, 20] 会在 7 之后返回 15，跳过了 9。
  • 错误写法：hasNext 用一个计数器与节点总数比较，却在构造时统计错了节点数 → 计数偏大时取完元素后仍返回真，next 对空栈弹出直接抛异常。
  • 错误写法：next 先压栈再弹栈，顺序颠倒 → 刚压入的右子树节点被当成本次答案，树 [7, 3, 15] 第一次调用就返回 15。
  • 错误写法：返回值取 stack.peek().val 却忘记弹出 → 同一个最小值被无限次返回，hasNext 永远为真。
  • 错误写法：Go 中方法用值接收者 func (this BSTIterator) Next() int → 栈的修改作用在副本上，每次 next 都从同一状态开始，反复返回同一个值。
  • 错误写法：构造函数只压入根节点而不压左链 → 树 [7, 3, 15] 第一次 next 返回 7 而不是 3，序列不再有序。
  • 错误写法：把 hasNext 写成"栈非空或当前指针非空"，却没有同步维护那个指针 → 指针残留导致取完元素后仍报告有下一个，随后弹空栈崩溃。
  • 错误写法：用递归中序配合"回调里返回"试图实现暂停 → 递归一旦返回进度就丢失，第二次调用只能从头再遍历，next 退化为 $O(n)$。

相似题目

题目 难度 考察点
94. 二叉树的中序遍历 简单 一次性输出完整中序序列，栈无需跨调用保留
230. 二叉搜索树中第 K 小的元素 中等 只需前 K 个元素，考察何时提前终止遍历
LCR 052. 递增顺序搜索树 简单 同样借中序推进，但要边遍历边改写指针
341. 扁平化嵌套列表迭代器 中等 嵌套结构的迭代器，栈中存的是迭代位置而非节点
281. 锯齿迭代器 中等 多路轮转取值，状态是队列轮次而非遍历路径
426. 将二叉搜索树转化为排序的双向链表 中等 一次性改成双向链表，此后前后移动都是 $O(1)$