【LeetCode 面试经典150题】230. Kth Smallest Element in a BST 二叉树中第k小的元素

230. Kth Smallest Element in a BST

题目大意

Given the root of a binary search tree, and an integer k, return the kth smallest value (1-indexed) of all the values of the nodes in the tree.

中文释义

给定一个二叉搜索树的根节点和一个整数 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）来查找第 k 小的元素。

1. 递归中序遍历DFS：

   • 实现 in_dfs 函数递归地进行中序遍历。
   • 中序遍历BST将以升序的方式访问所有节点。

2. 跟踪当前节点和索引：

   • 使用 current_num 来存储当前节点的值。
   • 使用 current_index 来跟踪当前遍历的节点索引。

3. 查找第 k 小的元素：

   • 当 current_index 达到 k 时，将 current_num 设置为答案 ans。

4. 返回结果：

   • 在 kthSmallest 函数中，初始化 current_index 和 ans，设置 k，然后调用 in_dfs 函数开始递归，最后返回 ans。

代码

class Solution {
public:
    int current_num;
    int current_index;
    int ans;
    int k;
    void in_dfs(TreeNode* root) {
        if (root && ans == INT_MAX) {
            in_dfs(root -> left);
            current_num = root -> val;
            current_index = current_index + 1;
            if (current_index == k) {
                ans = current_num;
            }
            in_dfs(root -> right);
        }
    }
    
    int kthSmallest(TreeNode* root, int k) {
        current_index = 0;
        ans = INT_MAX;
        this -> k = k;
        in_dfs(root);
        return ans;
    }
};