leetcode打卡--230. 二叉搜索树中第K小的元素

Posted 2021-11-04 C_YCBX Py_YYDS

文章目录

• • 题目
  • 题目解析
  • 解体代码
  • • 方法一：简单迭代
    • 方法二：记录size的结点
    • 方法三：转AVLTree

题目

题目描述

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

示例1：

示例2：

题目解析

由于是二叉搜索树，所以可以利用它的性质，直接得出中序遍历的按序元素的第k个即可，当然也可以用栈模拟迭代得出。

如果面试官加大难度，增加为以下两种拓展，你该怎么办？

1. 如果你需要频繁地查找第 k 小的值，你将如何优化算法？
2. 如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 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);
    }
};