230. 二叉搜索树中第 K 小的元素

力扣题目链接 (opens new window)

题目描述

给定一个二叉搜索树的根节点 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小的元素，就是中序遍历的第k个节点。

具体做法有两种：

方法一：中序遍历存数组

1. 中序遍历整棵树，把所有值存到数组里
2. 返回数组的第 k-1 个元素（因为数组下标从0开始）

这种方法简单直接，但要遍历整棵树，时间复杂度O(n)，空间复杂度O(n)。

方法二：中序遍历计数（推荐）

1. 中序遍历的过程中计数
2. 当计数到第k个节点时，直接返回，不需要遍历整棵树

这种方法平均情况下更快，如果k比较小，可以提前终止遍历。

下面用方法二实现。用一个全局变量count记录当前是第几个节点，用result记录答案。

模拟过程

以示例 2 为例，树的结构：

        5
       / \
      3   6
     / \
    2   4
   /
  1

中序遍历顺序：1 → 2 → 3 → 4 → 5 → 6

当 k = 3 时，遍历到节点 3 就是答案。

解题代码

class Solution {
private:
    int count = 0;
    int result = 0;
    
    void inorder(TreeNode* root, int k) {
        if (root == nullptr) return;
        
        inorder(root->left, k); // 左
        
        count++; // 中
        if (count == k) {
            result = root->val;
            return; // 找到答案，提前返回
        }
        
        inorder(root->right, k); // 右
    }
    
public:
    int kthSmallest(TreeNode* root, int k) {
        inorder(root, k);
        return result;
    }
};

复杂度分析

• 时间复杂度：O(H + k)，H 是树的高度，最坏情况O(n)
• 空间复杂度：O(H)，递归栈的深度

其他语言

Python3

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        self.count = 0
        self.result = 0
        
        def inorder(node):
            if not node:
                return
            
            inorder(node.left)  # 左
            
            self.count += 1  # 中
            if self.count == k:
                self.result = node.val
                return
            
            inorder(node.right)  # 右
        
        inorder(root)
        return self.result

Java

class Solution {
    private int count = 0;
    private int result = 0;
    
    public int kthSmallest(TreeNode root, int k) {
        inorder(root, k);
        return result;
    }
    
    private void inorder(TreeNode root, int k) {
        if (root == null) return;
        
        inorder(root.left, k); // 左
        
        count++; // 中
        if (count == k) {
            result = root.val;
            return;
        }
        
        inorder(root.right, k); // 右
    }
}

Go

func kthSmallest(root *TreeNode, k int) int {
    count := 0
    result := 0
    
    var inorder func(*TreeNode)
    inorder = func(node *TreeNode) {
        if node == nil {
            return
        }
        
        inorder(node.Left) // 左
        
        count++ // 中
        if count == k {
            result = node.Val
            return
        }
        
        inorder(node.Right) // 右
    }
    
    inorder(root)
    return result
}

JavaScript

var kthSmallest = function(root, k) {
    let count = 0;
    let result = 0;
    
    function inorder(node) {
        if (!node) return;
        
        inorder(node.left); // 左
        
        count++; // 中
        if (count === k) {
            result = node.val;
            return;
        }
        
        inorder(node.right); // 右
    }
    
    inorder(root);
    return result;
};

与代码随想录联系

本题是二叉搜索树的经典应用。在代码随想录二叉树章节中，详细讲解了：

• 二叉树的中序遍历
• 验证二叉搜索树

BST的核心性质：中序遍历得到升序序列。 这一点在很多题目中都会用到。

录友们如果对二叉树的遍历还不够熟练，建议系统学习 代码随想录二叉树章节，从递归三部曲开始，打好基础。

进阶思考

题目问：如果BST经常被修改，如何优化？

一种方法是给每个节点维护一个 leftCount 字段，表示左子树的节点数。这样查找第k小时：

• 如果 k <= leftCount，在左子树找第k小
• 如果 k == leftCount + 1，当前节点就是答案
• 如果 k > leftCount + 1，在右子树找第 k - leftCount - 1 小

这样查找时间复杂度O(H)，但每次插入/删除都要更新路径上所有节点的leftCount，维护成本较高。

平衡方案是用平衡二叉搜索树（如红黑树）+ leftCount字段，可以保证查找、插入、删除都是O(log n)。