从图书馆找书到二叉搜索树:如何快速找到第K小的元素
生活中的算法
想象你在图书馆的书架前,所有书籍都按照编号从小到大排列。现在要找到第5本编号最小的书,你会怎么做?很简单,只需要从左往右数到第5本即可。
这正是二叉搜索树(BST)的天然特性——中序遍历的结果就是有序序列。今天我们要解的这道题,就像在BST这座"数字图书馆"中,快速找到第K本"书"。
问题描述
LeetCode第230题"二叉搜索树中第K小的元素"要求:给定一个二叉搜索树的根节点和一个整数k,找出其中第k小的元素。
例如,给定BST:
3
/ \
1 4
\
2当k=1时,返回1;k=3时,返回3。
最直观的解法:中序遍历法
就像在图书馆里把书全部搬到地上再数到第k本,我们可以通过中序遍历将BST转换为有序数组,然后直接取第k-1个元素。
算法步骤
- 对BST进行中序遍历,得到升序数组
- 返回数组中第k-1个元素
用示例树模拟这个过程:
中序遍历顺序:1→2→3→4
当k=3时,数组第三个元素是3Java实现:
java
class Solution {
public int kthSmallest(TreeNode root, int k) {
List<Integer> list = new ArrayList<>();
inorder(root, list);
return list.get(k-1);
}
private void inorder(TreeNode node, List<Integer> list) {
if (node == null) return;
inorder(node.left, list);
list.add(node.val);
inorder(node.right, list);
}
}优化解法:遍历时提前返回
但如果我们只需要第k本,何必要把整个书架都搬空?发现数到第k本时可以直接停止!
改进思路
- 在中序遍历过程中记录访问次序
- 当计数器等于k时立即返回结果
- 利用BST特性提前终止遍历
算法步骤(迭代法)
- 使用栈模拟中序遍历
- 每次弹出节点时计数器加1
- 当计数器等于k时立即返回当前节点值
用示例树模拟(k=3):
栈操作流程:
1. 3入栈→3的左孩子1入栈→1的左孩子null
2. 弹出1,计数器=1≠3
3. 处理1的右子树2→2入栈→2的左孩子null
4. 弹出2,计数器=2≠3
5. 弹出3,计数器=3→找到答案3!Java代码:
java
class Solution {
public int kthSmallest(TreeNode root, int k) {
Deque<TreeNode> stack = new ArrayDeque<>();
TreeNode cur = root;
int count = 0;
while (cur != null || !stack.isEmpty()) {
// 深入最左端
while (cur != null) {
stack.push(cur);
cur = cur.left;
}
cur = stack.pop();
if (++count == k) return cur.val;
// 转向右子树
cur = cur.right;
}
return -1; // 不会执行到这里
}
}终极优化:Morris遍历法
如果连栈都不想用怎么办?就像在书架间穿梭时,用临时标记记录回程路线!
Morris遍历原理
- 利用叶子节点的空指针记录回溯路径
- 在遍历时临时修改树结构,之后恢复
- 实现O(1)空间复杂度
算法步骤
- 当前节点cur初始化为根节点
- 当cur不为空时循环:
- 如果cur无左子树:
- 计数器加1,若等于k则返回
- cur移向右子树
- 否则:
- 找到cur左子树的最右节点pre
- 若pre的右指针为空:将其指向cur,cur移向左子树
- 若pre的右指针为cur:恢复为空,处理当前节点,cur移向右子树
- 如果cur无左子树:
示例运行(k=3):
初始状态:
3
/ \
1 4
\
2
步骤:
1. cur=3,左子树存在
2. pre=2(1的右子树的最右)
3. pre.right=3,cur=1
4. cur=1,左子树不存在→计数器=1≠3→cur=1.right=2
5. cur=2,左子树不存在→计数器=2≠3→cur=2.right=3
6. 此时pre=2的右指针指向3→恢复pre.right=null→计数器+1=3→返回3Java实现:
java
class Solution {
public int kthSmallest(TreeNode root, int k) {
int count = 0;
TreeNode cur = root;
while (cur != null) {
if (cur.left == null) {
if (++count == k) return cur.val;
cur = cur.right;
} else {
// 找左子树的最右节点
TreeNode pre = cur.left;
while (pre.right != null && pre.right != cur) {
pre = pre.right;
}
if (pre.right == null) { // 建立线索
pre.right = cur;
cur = cur.left;
} else { // 拆除线索
pre.right = null;
if (++count == k) return cur.val;
cur = cur.right;
}
}
}
return -1;
}
}解法对比
| 方法 | 时间复杂度 | 空间复杂度 | 特点 |
|---|---|---|---|
| 递归中序遍历 | O(n) | O(n) | 简单但空间消耗大 |
| 迭代中序遍历 | O(n) | O(h) | 最优平衡(h为树高) |
| Morris遍历 | O(n) | O(1) | 空间最优但修改树结构 |
题目模式总结
这道题揭示了BST类问题的核心解法:
- 中序遍历有序性:BST问题的解题基石
- 遍历优化:通过提前终止或空间优化提升效率
- Morris技巧:在有限制条件下的空间优化方案
同类问题扩展:
- 验证二叉搜索树
- BST转换为累加树
- BST中的众数
解决这类问题的通用思路:
- 确认是否利用中序特性
- 根据空间限制选择遍历方式
- 在遍历过程中记录关键信息
小结
通过这道题,我们不仅掌握了三种不同时空复杂度的解法,更重要的是理解了BST的核心特性——中序遍历的有序性。就像在有序的书架上找书,关键是要掌握高效的检索方法。
记住:算法优化往往是在时空复杂度之间寻找平衡。在面试中,通常优先推荐迭代法中序遍历法,既保证了O(h)的空间复杂度,又易于理解和实现。
作者:忍者算法
公众号:忍者算法
我整理了LeetCode中最经典的100道算法题,并按照算法类型和难度等级分类。这些题目覆盖了90%以上的面试高频考点,公众号后台回复【百题清单】获取完整列表。刷完这些题,你会对算法有全新的认识!