Medium

题目描述

给定一个二叉搜索树的根节点 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 小的值，你将如何优化算法？

提示：

  • 尝试利用二叉搜索树的性质
  • 尝试中序遍历
  • 如果可以修改二叉搜索树节点的结构呢？
  • 最优的运行时间复杂度是 O(BST 的高度)

解题思路

解题思路

这道题要求在二叉搜索树中找到第 k 小的元素。由于 BST 的性质（左子树所有节点值 < 根节点值 < 右子树所有节点值），我们可以利用这个特性来解决问题。

方法一：中序遍历（推荐）

BST 的中序遍历会按照从小到大的顺序访问所有节点。因此，我们只需要进行中序遍历，当访问到第 k 个节点时，该节点的值就是答案。为了提高效率，我们可以在找到第 k 个节点后立即返回，而不需要遍历整棵树。

方法二：迭代中序遍历

使用栈来模拟递归的中序遍历过程，同样在访问第 k 个节点时返回结果。

方法三：分治法（最优）

利用 BST 的性质，我们可以通过统计左子树的节点数来判断第 k 小的元素在哪个子树中：

  • 如果左子树节点数 ≥ k，则第 k 小的元素在左子树
  • 如果左子树节点数 = k-1，则根节点就是第 k 小的元素
  • 否则，第 k 小的元素在右子树，且是右子树中第 (k - 左子树节点数 - 1) 小的元素

这里我们采用中序遍历的递归实现，因为代码简洁且易于理解。

代码实现

class Solution {
private:
    int count = 0;
    int result = 0;
    
    void inorder(TreeNode* root, int k) {
        if (!root) return;
        
        inorder(root->left, k);
        
        count++;
        if (count == k) {
            result = root->val;
            return;
        }
        
        inorder(root->right, k);
    }
    
public:
    int kthSmallest(TreeNode* root, int k) {
        count = 0;
        inorder(root, k);
        return result;
    }
};

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        def inorder(node):
            if not node:
                return []
            return inorder(node.left) + [node.val] + inorder(node.right)
        
        # 简化版本：直接返回第k-1个元素
        return inorder(root)[k-1]
        
        # 优化版本（提前返回）：
        # self.count = 0
        # self.result = 0
        
        # def inorder_optimized(node):
        #     if not node:
        #         return
            
        #     inorder_optimized(node.left)
            
        #     self.count += 1
        #     if self.count == k:
        #         self.result = node.val
        #         return
            
        #     inorder_optimized(node.right)
        
        # inorder_optimized(root)
        # return self.result

public class Solution {
    private int count = 0;
    private int result = 0;
    
    public int KthSmallest(TreeNode root, int k) {
        count = 0;
        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);
    }
}

var kthSmallest = function(root, k) {
    let count = 0;
    let result = 0;
    
    function inorder(node) {
        if (!node || count >= k) return;
        
        inorder(node.left);
        
        count++;
        if (count === k) {
            result = node.val;
            return;
        }
        
        inorder(node.right);
    }
    
    inorder(root);
    return result;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

算法时间复杂度空间复杂度
中序遍历（递归）O(H + k)O(H)
中序遍历（迭代）O(H + k)O(H)

其中 H 是树的高度。在最坏情况下（树退化为链表），H = n；在最好情况下（平衡二叉树），H = log n。

注意： 这里的时间复杂度是 O(H + k) 而不是 O(n)，因为我们在找到第 k 个元素后就停止遍历了。

相关题目