给一棵二叉搜索树,写一个 KthSmallest 函数来找到其中第 K 小的元素。
样例
样例 1:
输入:{1,#,2},2
输出:2
解释:
1
\
2
第二小的元素是2。
样例 2:
输入:{2,1,3},1
输出:1
解释:
2
/ \
1 3
第一小的元素是1。
挑战
如果这棵 BST 经常会被修改(插入/删除操作)并且你需要很快速的找到第 K 小的元素呢?你会如何优化这个 KthSmallest 函数?
注意事项
你可以假设 k 总是有效的, 1 ≤ k ≤ 树的总节点数。
/**
* Definition of TreeNode:
* class TreeNode {
* public:
* int val;
* TreeNode *left, *right;
* TreeNode(int val) {
* this->val = val;
* this->left = this->right = NULL;
* }
* }
*/
class Solution {
public:
/**
* @param root: the given BST
* @param k: the given k
* @return: the kth smallest element in BST
*/
int kthSmallest(TreeNode * root, int k) {
// write your code here
vector<int> res;
kthSmallest(root,res);
return res[k-1];
}
void kthSmallest(TreeNode * root,vector<int> &res) {
// write your code here
if(root==NULL) return;
kthSmallest(root->left,res);
res.push_back(root->val);
kthSmallest(root->right,res);
}
};
/**
* Definition of TreeNode:
* class TreeNode {
* public:
* int val;
* TreeNode *left, *right;
* TreeNode(int val) {
* this->val = val;
* this->left = this->right = NULL;
* }
* }
*/
class Solution {
public:
/**
* @param root: the given BST
* @param k: the given k
* @return: the kth smallest element in BST
*/
int kthSmallest(TreeNode * root, int k) {
// write your code here
int res=0;
kthSmallest(root,k,res);
return res;
}
void kthSmallest(TreeNode * root,int &k ,int&res) {
// write your code here
if(root==NULL) return;
kthSmallest(root->left,k,res);
k--;
if(k==0) {
res=root->val;
}
kthSmallest(root->right,k,res);
}
};