【LeetCode热题100道笔记】二叉搜索树中第 K 小的元素

题目描述

给定一个二叉搜索树的根节点 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<=10^4$
$0<=Node.val<=10^4$

进阶：如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

思考一：中序遍历（利用BST有序特性）

核心是借助BST中序遍历的有序性：中序遍历（左→根→右）会生成严格递增的节点值序列，序列中第k-1个元素（索引从0开始）即为第k小元素。

算法过程

1. 中序遍历收集节点值：
   • 初始化空数组arr，用于存储中序遍历的节点值；
   • 递归执行中序遍历：先遍历左子树（确保更小的元素先入数组），再将当前节点值加入arr，最后遍历右子树；
   • 若节点为空，直接返回（递归终止条件）。
2. 提取第k小元素：
   • 遍历完成后，arr已按递增排序，返回arr[k-1]（第k个元素，索引为k-1）。

时空复杂度

• 时间复杂度：O(n)，n为二叉树节点总数。
  原因：中序遍历需遍历所有节点（最坏情况需收集全部节点值才能找到第k小），总操作次数与节点数线性相关。
• 空间复杂度：O(n)。
  原因：数组arr需存储所有节点值（最坏情况），递归调用栈需O(h)（h为树高），总空间由数组主导，为O(n)。

代码

/*** Definition for a binary tree node.* function TreeNode(val, left, right) {*     this.val = (val===undefined ? 0 : val)*     this.left = (left===undefined ? null : left)*     this.right = (right===undefined ? null : right)* }*/
/*** @param {TreeNode} root* @param {number} k* @return {number}*/
var kthSmallest = function(root, k) {const arr = [];inOrder(root, arr);return arr[k-1];
};function inOrder(node, arr) {if (!node) return;inOrder(node.left, arr);arr.push(node.val);inOrder(node.right, arr);
}

思考二：最大堆（动态维护最小的k个元素）

核心是用容量为k的最大堆动态筛选“当前遍历到的最小k个元素”：堆顶始终是这k个元素中的最大值，遍历完成后堆顶即为第k小元素（因比堆顶小的元素有k-1个，堆顶自然是第k小）。
注：JavaScript无内置堆，需实现简易最大堆（按值大小排序，堆顶为最大值）。

算法过程

1. 实现最大堆：
   • 堆初始化时固定容量为k，超过容量时自动弹出堆顶（最大值）；
   • 提供push（插入元素并调整堆结构）、front（获取堆顶元素）、size（获取堆元素个数）方法。
2. 深度优先遍历（DFS）树节点：
   • 遍历所有节点，对每个节点值：
     • 若堆中元素不足k个，直接入堆；
     • 若堆已满且当前节点值 < 堆顶（说明当前节点值是更小的元素，需替换堆顶），弹出堆顶后将当前节点值入堆；
3. 返回结果：遍历完成后，堆顶元素即为第k小元素。

时空复杂度

• 时间复杂度：O(n log k)，n为二叉树节点总数。
  原因：遍历所有节点需O(n)，每个节点入堆/出堆操作需O(log k)（堆的调整时间与堆容量k的对数相关），总时间为O(n log k)。
• 空间复杂度：O(k + h)，h为树高。
  原因：堆占用O(k)空间，DFS递归栈占用O(h)空间，总空间由堆主导（k远小于n时比中序遍历更优）。

代码

/*** Definition for a binary tree node.* function TreeNode(val, left, right) {*     this.val = (val===undefined ? 0 : val)*     this.left = (left===undefined ? null : left)*     this.right = (right===undefined ? null : right)* }*/
/*** @param {TreeNode} root* @param {number} k* @return {number}*/
var kthSmallest = function(root, k) {const priorityQueue = new MyMaxPriorityQueue(k);dfs(root, priorityQueue, k);return priorityQueue.front();    
};function dfs(node, queue, k) {if (!node) return;if (queue.size() < k || queue.front() > node.val) {queue.push(node.val);}dfs(node.left, queue, k);dfs(node.right, queue, k);
}class MyMaxPriorityQueue {constructor(capacity = 1000) {this._data = [];this._capacity = capacity;this._size = 0;}front() {return this._data[0];}push(num) {if (this._capacity === this._size) {this.pop();}this._data.push(num);this.swim();this._size++;}pop() {if (this._data.length === 0) return;[this._data[0], this._data[this._data.length-1]] = [this._data[this._data.length-1], this._data[0]];const item = this._data.pop();this.sink();this._size--;return item;}swim(index = this._data.length-1) {while (index > 0) {let pIndex = Math.floor((index-1)/2);if (this._data[index] > this._data[pIndex]) {[this._data[index], this._data[pIndex]] = [this._data[pIndex], this._data[index]];index = pIndex;continue;}break;}}sink(index = 0) {const n = this._data.length;while (true) {let left = 2 * index + 1;let right = left + 1;let biggest = index;if (left < n && this._data[left] > this._data[index]) {biggest = left;}if (right < n && this._data[right] > this._data[biggest]) {biggest = right;}if (biggest !== index) {[this._data[biggest], this._data[index]] = [this._data[index], this._data[biggest]];index = biggest;continue;}break;}}size() {return this._size;}}