题目描述
给定一个二叉搜索树的根节点
root
,和一个整数k
,请你设计一个算法查找其中第k
小的元素(从 1 开始计数)。示例 1:
输入:root = [3,1,4,null,2], k = 1 输出:1示例 2:
输入:root = [5,3,6,2,4,null,null,1], k = 3 输出:3提示:
- 树中的节点数为
n
。1 <= k <= n <= 104
0 <= Node.val <= 104
进阶:如果二叉搜索树经常被修改(插入/删除操作)并且你需要频繁地查找第
k
小的值,你将如何优化算法?
解题思路
要查找二叉搜索树(BST)中的第 k
小元素,我们可以利用 BST 的中序遍历特性。中序遍历 BST 会以升序方式访问所有节点,因此第 k
小的元素就是中序遍历结果中的第 k
个元素。
-
中序遍历:中序遍历 BST 的结果是一个升序排列的节点值列表。可以使用递归或迭代的方式进行中序遍历。
-
递归方法:在遍历过程中,维护一个计数器来记录已经遍历的节点数量,当计数器等于
k
时,返回当前节点的值。
复杂度分析
- 时间复杂度:
O(n)
,其中n
是树的节点数。需要遍历整个树的节点,直到找到第k
小的元素。 - 空间复杂度:
O(h)
,其中h
是树的高度。递归栈的空间复杂度为树的高度。
代码实现
package org.zyf.javabasic.letcode.hot100.tree;
import org.zyf.javabasic.letcode.tree.base.TreeNode;
/**
* @program: zyfboot-javabasic
* @description: 二叉搜索树中第 K 小的元素(中等)
* @author: zhangyanfeng
* @create: 2024-08-22 11:48
**/
public class KthSmallestSolution {
private int count = 0; // Count of nodes visited
private int result = -1; // To store the kth smallest value
public int kthSmallest(TreeNode root, int k) {
// Perform in-order traversal
inOrderTraversal(root, k);
return result;
}
private void inOrderTraversal(TreeNode node, int k) {
if (node == null) {
return;
}
// Traverse left subtree
inOrderTraversal(node.left, k);
// Visit node
count++;
if (count == k) {
result = node.val;
return;
}
// Traverse right subtree
inOrderTraversal(node.right, k);
}
public static void main(String[] args) {
KthSmallestSolution solution = new KthSmallestSolution();
// Example 1
TreeNode root1 = new TreeNode(3);
root1.left = new TreeNode(1);
root1.right = new TreeNode(4);
root1.left.right = new TreeNode(2);
System.out.println(solution.kthSmallest(root1, 1)); // Output: 1
// Example 2
TreeNode root2 = new TreeNode(5);
root2.left = new TreeNode(3);
root2.right = new TreeNode(6);
root2.left.left = new TreeNode(2);
root2.left.right = new TreeNode(4);
root2.left.left.left = new TreeNode(1);
System.out.println(solution.kthSmallest(root2, 3)); // Output: 3
}
}