44 二叉搜索树中第K小的元素

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 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、将根节点和左子树的节点入栈

cpp 复制代码

stack.push(root);

2、一直走向他的左子树，循环入栈，直到节点为空

cpp 复制代码

while(root){
    root=root->left;
    stack.push(root);
}

3、栈顶出栈，走向右子树节点，出栈第k次即为第k小的元素，直接返回

cpp 复制代码

root=stack.top();
root=root->right;

4、第二次循环，继续持续走root节点的左子树，让左子树入栈

cpp 复制代码

while(root){
    root=root->left;
    stack.push(root);
}

5、栈为空和root为空时退出循环。

由于题目的限制，不需要考虑k不满足条件的情况。

cpp 复制代码

int kthSmallest(TreeNode* root, int k) {
    if(!root)   return -3;
    stack<TreeNode*> stack;
    while(root||!stack.empty()){
        //一直往左走
        while(root){
            stack.push(root);
            root=root->left;
        }
        //root为空时出栈
        if(!stack.empty()){
            root=stack.top();
            stack.pop();
            k--;
            if(0==k) return root->val;
        }
        root=root->right;
    }
    return 0;   
}

如果你需要频繁地查找第 k 小的值，你将如何优化算法？

可以记录下以每个结点为根结点的子树的结点数，并在查找第 k 小的值时，使用如下方法搜索：令 node 等于根结点，开始搜索。

对当前结点 node 进行如下操作：

【1】如果 node 的左子树的结点数 left 小于 k−1，则第 k 小的元素一定在 node 的右子树中，令 node 等于其的右子结点，k 等于 k−left−1，并继续搜索。

【2】如果 node 的左子树的结点数 left 等于 k−1，则第 k 小的元素即为 node ，结束搜索并返回 node 即可。

【3】如果 node 的左子树的结点数 left 大于 k−1，则第 k 小的元素一定在 node 的左子树中，令 node 等于其左子结点，并继续搜索。

cpp 复制代码

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);
    }
};

推荐一个零声教育学习教程，个人觉得老师讲得不错，分享给大家：[Linux，Nginx ，ZeroMQ，MySQL，Redis，fastdfs，MongoDB，ZK，流媒体，CDN，P2P，K8S，Docker，TCP/IP，协程，DPDK等技术内容，点击立即学习:链接