【中等】力扣算法题解析LeetCode230：二叉搜索树中第K小的元素

关注文末推广名片，即可免费获得本题测试源码！

题目来源：LeetCode230：二叉搜索树中第K小的元素

问题抽象： 给定一个二叉搜索树（BST）的根节点 root 和整数 k，要求找出 BST 中 第 k 小的元素值（按升序排列的第 k 个元素），需满足以下核心需求：

1. BST 特性利用：

   • 利用 BST 中序遍历升序 的特性（左子树值 < 根值 < 右子树值）；
   • 无需完全展开树结构（避免 O(n) 空间复杂度）。

2. 输入约束：

   • 节点数 n ∈ [1, 10^4]；
   • k ∈ [1, n]（保证有效索引）；
   • 节点值唯一（含负整数）。

3. 计算要求：

   • 时间复杂度 O(h + k)（h 为树高，最坏 O(n) 当 k=n）；
   • 空间复杂度 O(h)（递归栈深度，最坏 O(n) 当树退化为链表）。

4. 关键策略：

   • 中序遍历优化：
     1. 递归或迭代遍历左子树；
     2. 访问根节点时计数 +1；
     3. 当计数 =k 时立即返回当前节点值（提前终止遍历）；
     4. 否则遍历右子树。

5. 边界处理：

   • k=1 时返回最小值（最左下叶子节点）；
   • k=n 时返回最大值（最右下叶子节点）；

输入：BST 根节点 root；整数 k（第 k 小元素的索引）。
输出：节点值（第 k 小的元素值）。

解题思路

1. 中序遍历特性：
   二叉搜索树（BST）的中序遍历结果是一个升序序列，因此第 k 小的元素即为中序遍历序列中的第 k 个元素。

2. 迭代优化：
   使用迭代方式实现中序遍历，在遍历过程中计数，当计数达到 k 时立即返回当前节点值，避免完整遍历整棵树，降低时间复杂度。

3. 算法步骤：

   • 初始化一个栈用于模拟递归过程。
   • 从根节点开始，将所有左子节点依次入栈（形成左侧链）。
   • 弹出栈顶节点（当前最小节点），计数器减一：
     • 若计数器为 0，返回当前节点值。
     • 否则，将当前节点的右子节点作为新的根节点，重复上述过程（处理右子树）。

代码实现（Java版）🔥点击下载源码

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        // 使用栈模拟中序遍历过程
        Deque<TreeNode> stack = new ArrayDeque<>();
        TreeNode current = root;
        
        while (current != null || !stack.isEmpty()) {
            // 遍历到最左侧节点，依次入栈
            while (current != null) {
                stack.push(current);
                current = current.left;
            }
            
            // 弹出当前最小节点
            current = stack.pop();
            // 每弹出一个节点，k减1
            if (--k == 0) {
                return current.val; // 找到第k小元素
            }
            
            // 转向右子树
            current = current.right;
        }
        
        return -1; // 未找到（题目保证k有效，此处不会执行）
    }
}

代码说明

1. 栈模拟中序遍历：

   • 使用 Deque 实现栈（推荐替代 Stack）。
   • 通过循环将左子节点全部入栈，确保每次弹出当前最小节点。

2. 提前终止：

   • 每次弹出节点时，k 减 1，当 k 减至 0 时立即返回结果，避免无效遍历。

3. 时间复杂度：

   • 最优：当 k 较小时（如 k=1），时间复杂度为 O(H)（H 为树高）。
   • 最坏：当 k=n 时，需遍历所有节点，时间复杂度为 O(n)。

4. 空间复杂度：

   • 栈空间最多存储 H 个节点，空间复杂度为 O(H)（树高）。

进阶优化思路

若需频繁修改树结构并查询第 k 小值，可考虑以下优化：

1. 维护子树节点数：

   • 在每个节点中增加字段 size，记录以该节点为根的子树节点总数。
   • 插入/删除时更新路径上的 size 值（时间复杂度 O(H)）。

2. 快速查询算法：

   • 从根节点开始，设左子树节点数为 leftSize：
     • 若 k <= leftSize，在左子树中查找第 k 小。
     • 若 k == leftSize + 1，当前节点即为目标。
     • 若 k > leftSize + 1，在右子树中查找第 k - leftSize - 1 小。
   • 查询时间复杂度降至 O(H)。

3. 平衡二叉搜索树：

   • 使用 AVL 树或红黑树保证树高 H = O(log n)，使插入、删除、查询操作均稳定在 O(log n)。

提交详情（执行用时、内存消耗）