230. 二叉搜索树中第 K 小的元素 - 力扣(LeetCode)
一、题目要求
给定一个二叉搜索树的根节点 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 <= 104
0 <= Node.val <= 104
进阶:如果二叉搜索树经常被修改(插入/删除操作)并且你需要频繁地查找第 k
小的值,你将如何优化算法?
二、解法1-通用做法 O(N)
解法1对任意一颗二叉树都能找到第K小的元素,不论是否是二叉搜索树。
class Solution {
public:int kthSmallest(TreeNode* root, int k) {if(root == nullptr || k < 0) // 安全检查return -1;map<int,int> map;queue<TreeNode*> queue;queue.push(root);while(!queue.empty()){root = queue.front();map[root->val]++;if(root->left)queue.push(root->left);if(root->right)queue.push(root->right);queue.pop();}for(const auto& it:map){k-=it.second;if(k <= 0)return it.first;}return -1; // k比节点数还大,则不存在}
};
优化1:
二叉搜索树的一个特征是没有重复的元素,我们可以利用这个特点,把所有的节点值插入一个数组容器中,排序后下标为k-1的元素就是第k小的元素。
class Solution {
public:int kthSmallest(TreeNode* root, int k) {if(root == nullptr || k < 0) // 安全检查return -1;vector<int> v;queue<TreeNode*> queue;queue.push(root);while(!queue.empty()){root = queue.front();v.push_back(root->val);if(root->left)queue.push(root->left);if(root->right)queue.push(root->right);queue.pop();} sort(v.begin(),v.end()); return v[k-1];}
};
优化2:
二叉搜索树最重要的特性是左子树的所有节点比根小,右子树的所有节点比根大,那么它的中序遍历一定是一个升序的序列,就不需要额外排序了,将上述的插入顺序改为中序遍历。
class Solution {
public:int kthSmallest(TreeNode* root, int k) {if (root == nullptr || k < 0) // 安全检查return -1;stack<TreeNode*> stack;int ret;for(int i = 1;i <= k;i++) { // 只需要得到第k个即可while (root) {stack.push(root);root = root->left;}root = stack.top()->right;ret = stack.top()->val;stack.pop();}return ret;}
};