[数据结构与算法]【数据结构基础_树】Leetcode 230.二叉搜索树中第k小的元素

原题链接：Leetcode 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.

Example 1:

Input: root = [3,1,4,null,2], k = 1
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
Output: 3

Constraints:

• The number of nodes in the tree is n.
• 1 <= k <= n <= 10⁴
• 0 <= Node.val <= 10⁴

Follow up:

• If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?

方法一：中序遍历+数组

思路：

BST的特性是中序遍历就是增序的
那么可以通过中序遍历，将关键字存在数组中，然后输出第k个元素即可

C++

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */

class Solution {
public:
    // 递归实现中序遍历
    void dfs(TreeNode* root,vector<int>& nums){
        if(!root)
            return;

        dfs(root->left, nums);
        nums.push_back(root->val);
        dfs(root->right, nums);
    }

    int kthSmallest(TreeNode* root, int k) {
        vector<int> nums;
        dfs(root, nums);

        // 输出第k个小的数
        return nums[k-1];
    }
};

复杂度分析：

• 时间复杂度：O(n)，需要遍历所有元素一次
• 空间复杂度：O(n)，数据需要存所有元素

方法二：中序遍历+优先队列

思路：

方法1的局限性在于：
需要输出所有元素，存入数组，当数据量比较大的时候，时间、空间的消耗都比较大。
但其实只需要维护前k个元素就可以了。

优先队列，是解决第k大/小的数这类问题的首选之一。我们只需要在每次插入的时候判断队列长度是否大于k，如果大于则弹出队列中最大的元素。

C++代码：

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */

class Solution {
public:
    void dfs(TreeNode* root, priority_queue<int>& pq, int k){
        if(!root){
            return;
        }
        dfs(root->left, pq, k);
        pq.push(root->val);
        // 当元素个数大于k时，将最大元素弹出
        if(pq.size() > k){
            pq.pop();
        }
        dfs(root->right, pq, k);
    }

    int kthSmallest(TreeNode* root, int k) {
        // 优先队列
        priority_queue<int> pq;
        dfs(root, pq, k);
        return pq.top();
    }
};