Problem: 230. 二叉搜索树中第 K 小的元素
给定一个二叉搜索树的根节点 root ,和一个整数 k ,请你设计一个算法查找其中第 k 小的元素(从 1 开始计数)。
整体思路
这段代码旨在解决一个经典的二叉搜索树(BST)问题:二叉搜索树中第K小的元素 (Kth Smallest Element in a BST)。其目标是在一个有效的BST中,找到按升序排列的第 k
个节点的值。
该算法巧妙地利用了BST的一个核心性质:BST的中序遍历结果是一个严格递增的有序序列。因此,寻找第 k
小的元素就等价于寻找中序遍历序列中的第 k
个元素。
算法的核心思路是执行一次受控的中序遍历:
-
利用中序遍历的顺序:
- 算法通过一个递归的
dfs
函数来实现深度优先搜索,其结构dfs(left) -> process(node) -> dfs(right)
正是中序遍历的模式。 - 这意味着节点被处理的顺序(即
process(node)
被执行的顺序)是严格按照节点值从小到大的顺序。
- 算法通过一个递归的
-
计数与提前终止:
- 为了找到第
k
个元素,算法使用一个计数器cnt
。它被初始化为k
。 - 在(隐式的)中序遍历过程中,每当处理一个节点(即从左子树返回后),就将
cnt
减一。 - 当
cnt
减到0
时,说明当前处理的这个节点就是我们正在寻找的第k
小的元素。 - 此时,将其值存入结果变量
ans
中。
- 为了找到第
-
剪枝(隐式):
- 虽然代码结构上看起来会遍历整棵树,但一旦
ans
被赋值,后续的dfs
调用虽然会继续执行,但它们不会再改变ans
的值。 - 一个更优化的版本可以在找到答案后通过抛出异常或返回特殊标志来提前终止整个递归过程,但这会使代码稍微复杂一些。当前版本的实现虽然不提前终止,但逻辑正确且易于理解。
- 虽然代码结构上看起来会遍历整棵树,但一旦
-
状态管理:
cnt
和ans
被声明为类成员变量,以便在dfs
的递归调用之间共享和修改它们的状态。kthSmallest
方法负责初始化cnt
并启动递归,最后返回ans
。
完整代码
/**
* Definition for a binary tree node.
* public class TreeNode {
* int val;
* TreeNode left;
* TreeNode right;
* TreeNode() {}
* TreeNode(int val) { this.val = val; }
* TreeNode(int val, TreeNode left, TreeNode right) {
* this.val = val;
* this.left = left;
* this.right = right;
* }
* }
*/
class Solution {
// cnt: 倒数计数器,从 k 开始递减。
private int cnt;
// ans: 用于存储最终找到的第 k 小的元素的值。
private int ans;
/**
* 在二叉搜索树中查找第 k 小的元素。
* @param root BST的根节点
* @param k 要查找的元素的排名(从1开始)
* @return 第 k 小的元素的值
*/
public int kthSmallest(TreeNode root, int k) {
// 初始化计数器为 k
this.cnt = k;
// 启动中序遍历的递归过程
dfs(root);
// 返回在 dfs 过程中找到的结果
return this.ans;
}
/**
* 递归辅助函数,执行受控的中序遍历。
* @param node 当前访问的节点
*/
private void dfs(TreeNode node) {
// 基本情况:如果节点为空,则返回。
if (node == null) {
return;
}
// 1. 递归地访问左子树
dfs(node.left);
// --- 中序遍历的核心处理逻辑 ---
// 2. 处理当前节点
// 每处理一个节点,计数器减一
this.cnt--;
// 检查是否找到了第 k 个元素
if (this.cnt == 0) {
// 如果 cnt 变为 0,说明当前节点就是第 k 小的元素
this.ans = node.val;
// 注意:虽然找到了答案,但此版本的递归不会立即停止,
// 而是会完成右子树的空转。
// 可以在此处 return 来做一个简单的剪枝。
return;
}
// 3. 递归地访问右子树
// 如果已经找到答案 (cnt <= 0),可以避免进入右子树来优化
if (cnt > 0) {
dfs(node.right);
}
}
}
时空复杂度
时间复杂度:O(H + k)
- 节点访问:算法通过中序遍历的方式访问节点。它会首先深入到最左边的叶子节点,然后开始回溯并向右访问。
- 访问路径:
- 为了找到第
k
小的元素,算法需要完整地遍历包含这k
个元素的左侧子树。 - 在最坏的情况下(例如,
k=N
),算法需要遍历整棵树,时间复杂度为 O(N)。 - 在一般情况下,算法会访问从根节点到第
k
小元素路径上的所有节点,以及这些路径节点的部分子树。 - 具体来说,它会访问到第
k
小的元素为止。这个过程涉及的节点数大致等于从根到第k
小元素的路径长度(最坏为H
,树的高度)加上k
。 - 因此,一个更精确的时间复杂度是 O(H + k)。
- 为了找到第
- 特殊情况:
- 如果树是平衡的,
H
约等于log N
,复杂度为 O(log N + k)。 - 如果树是退化的链表,
H
约等于N
,复杂度为 O(N + k) = O(N)。 - 由于
k
最大可以是N
,所以最坏时间复杂度是 O(N)。
- 如果树是平衡的,
空间复杂度:O(H)
- 主要存储开销:算法的额外空间主要由递归调用栈所占用。
- 递归深度:递归的最大深度取决于树的高度
H
。- 最好情况:如果树是高度平衡的,其高度
H
约为log N
。此时空间复杂度为 O(log N)。 - 最坏情况:如果树是极度不平衡的(例如,退化成一个链表),其高度
H
将等于节点数N
。此时空间复杂度为 O(N)。
- 最好情况:如果树是高度平衡的,其高度
- 类成员变量:
cnt
和ans
占用的空间是 O(1)。
综合分析:
该算法的(辅助)空间复杂度由递归栈的深度决定,因此为 O(H),其中 H
是树的高度。