力扣 230. 二叉搜索树中第K小的元素

230. 二叉搜索树中第K小的元素#

给定一个二叉搜索树的根节点 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 小的值，你将如何优化算法？

未进阶#

不考虑进阶的问题，因为二叉搜索树的中序遍历得到的结果是升序的，可以通过递归或者迭代，中序遍历到第k个节点，然后返回。这里选择Morris中序遍历。

查看代码

 /**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int res=0;
        int cnt=0;//计数
        while (root != NULL)
        {
            if (root->left == NULL)
            {
                // cout << root->val << " ";//输出当前节点
                cnt++;
                if(cnt==k)
                    res=root->val;
                root = root->right;
            }
            else
            {
                // 找当前节点的前趋结点
                TreeNode* predecessor = root->left;
                while (predecessor->right != NULL
                    && predecessor->right != root)
                {
                    predecessor = predecessor->right;
                }

                // 使当前节点成为inorder的前序节点的右侧子节点
                if (predecessor->right == NULL)
                {
                    predecessor->right = root;
                    root = root->left;
                }
                //复原之前的修改
                else
                {
                    predecessor->right = NULL;
                    // cout << root->val << " ";//输出当前节点
                    cnt++;
                    if(cnt==k)
                        res=root->val;
                    root = root->right;
                }
            }
        }
        return res;
    }
};

进阶#

如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

记录子树节点#

既然频繁地查找第 k 小的值，那么就记录每个节点中序遍历前面有几个节点，记为pre_nums，在查找的时候，提高比较k和pre_nums的值，省略一部分搜索，提高查找效率。

• 如果k-1（因为在比较cur左节点）<cur左节点的pre_nums,则当前节点cur右子树就不需要去遍历了，k会在左子树中。
• 如果k-1>cur左节点的pre_nums,则当前节点cur左子树就不需要去遍历了，k会在右子树中。
• 如果k-1==cur左节点的pre_nums,则当前节点cur的pre_nums==k。

可以将pre_nums存储在结点中，也可以将其记录在哈希表中。

代码来自官方

查看代码

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);//本身（1）+左子树+右子树
        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);
    }
};

平衡二叉树#

等后面做到平衡二叉树再来补