题目
题目描述
给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。
示例1：
示例2：
题目解析
由于是二叉搜索树，所以可以利用它的性质，直接得出中序遍历的按序元素的第k个即可，当然也可以用栈模拟迭代得出。
如果面试官加大难度，增加为以下两种拓展，你该怎么办？
- 如果你需要频繁地查找第
k小的值，你将如何优化算法？ - 如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第
k小的值，你将如何优化算法？
这两个提问方式，肯定就是考察你对其他二叉搜索树的衍生数据结构的掌握。
第一个优化方式有很多种，你可以使用最简单的，直接就是把每个结点增加一个存储该树的结点个数的变量。
第二个优化方式，你可以选择AVL或者红黑树，当然AVL相对红黑树要更好写，所以我们还是选择了AVL。
解体代码
方法一：简单迭代
class Solution {
public:
int kthSmallest(TreeNode* root, int k) {
stack<TreeNode*>St;
//一直push左子树
while(root){
//得出所有的最小子树，根据栈的特性，则会从小到大排列
St.push(root);
root = root->left;
}
//开始中序模拟
while(!St.empty()){
TreeNode* t = St.top();
St.pop();
if(!(--k)){
return t->val;
}
//一旦右子树不为空，则继续push
if(t->right){
St.push(t->right);
//一旦右子树的左子树不为空，继续push
if(t->right->left){
t = t->right->left;
while(t){
St.push(t);
t = t->left;
}
}
}
}
return -1;
}
};
方法二：记录size的结点
class MyBst {
public:
MyBst(TreeNode *root) {
this->root = root;
countNodeNum(root);
}
// 返回二叉搜索树中第k小的元素
int kthSmallest(int k) {
TreeNode *node = root;
while (node != nullptr) {
int left = getNodeNum(node->left);
if (left < k - 1) {
node = node->right;
k -= left + 1;
} else if (left == k - 1) {
break;
} else {
node = node->left;
}
}
return node->val;
}
private:
TreeNode *root;
unordered_map<TreeNode *, int> nodeNum;
// 统计以node为根结点的子树的结点数
int countNodeNum(TreeNode * node) {
if (node == nullptr) {
return 0;
}
nodeNum[node] = 1 + countNodeNum(node->left) + countNodeNum(node->right);
return nodeNum[node];
}
// 获取以node为根结点的子树的结点数
int getNodeNum(TreeNode * node) {
if (node != nullptr && nodeNum.count(node)) {
return nodeNum[node];
}else{
return 0;
}
}
};
class Solution {
public:
int kthSmallest(TreeNode* root, int k) {
MyBst bst(root);
return bst.kthSmallest(k);
}
};
方法三：转AVLTree
class AVLTree{
TreeNode* root;
unordered_map<TreeNode*,int>check_height;
public:
AVLTree(TreeNode* root):root(root){
update_height(root);
root = convToAVL(root);
}
public:
//直接转AVL，转之前先记录好每一颗树的高度，由于不再好重新设计数据的结点结构，所以直接用哈希表记录每棵树的高度
int update_height(TreeNode* root){
if(root==nullptr) return 0;
int Llen = update_height(root->left)+1;
int Rlen = update_height(root->right)+1;
return check_height[root] = Llen>Rlen?Llen:Rlen;
}
int get_height(TreeNode* root){
if(root==nullptr)
return 0;
return check_height[root];
}
void update(TreeNode* root){
check_height[root] = get_height(root->left)>get_height(root->right)?get_height(root->left):get_height(root->right);
}
TreeNode* rotateLeft(TreeNode* root){
TreeNode* son = root->right;
root->right = son->left;
son->left = root;
update(root);
update(son);
return son;
}
TreeNode* rotateRight(TreeNode* root){
TreeNode* son = root->left;
root->left = son->right;
son->right = root;
update(root);
update(son);
return son;
}
TreeNode* rotateLeftRight(TreeNode* root){
root->left = rotateLeft(root->left);
return rotateLeft(root);
}
TreeNode* rotateRightLeft(TreeNode* root){
root->right = rotateRight(root->right);
return rotateLeft(root);
}
TreeNode* convToAVL(TreeNode* root){
if(root==nullptr)
return 0;
root->left = convToAVL(root->left);
root->right = convToAVL(root->right);
if(get_height(root->left)-get_height(root->right)==2){
root = get_height(root->left->left)>get_height(root->left->right)?rotateRight(root):rotateLeftRight(root);
}else if(get_height(root->left)-get_height(root->right)==-2){
root = get_height(root->right->right)>get_height(root->right->left)?rotateLeft(root):rotateRightLeft(root);
}
update(root);
return root;
}
int kth(int k){
stack<TreeNode *> stack;
while (root != nullptr || stack.size() > 0) {
while (root != nullptr) {
stack.push(root);
root = root->left;
}
root = stack.top();
stack.pop();
--k;
if (k == 0) {
break;
}
root = root->right;
}
return root->val;
}
};
class Solution {
public:
int kthSmallest(TreeNode* root, int k) {
AVLTree s(root);
return s.kth(k);
}
};