【每日一题】力扣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 <= 10⁴
• 0 <= Node.val <= 10⁴

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

解题思路

中序遍历

由于题目给的二叉树是二叉搜索树，所以可以根据中序遍历，把节点值存入数组中再直接返回第 K 大的数。这种方法比较简单，官方给的第二种方法比较麻烦，我就说一下大体的思路吧。第三种方法我是真的看不下去，直接手撕平衡树了，有兴趣的可以自己去看看。

优化

对于频繁的查找是可以优化的。

因为树是二叉树，所以想确定第 k 小的数的位置，可以根据树的子节点数确定第 K 小的数。这里用一个哈希表存储每个根节点的子节点数，这样可以遍历二叉树的节点查找第 k 小的数。

• 如果节点的左子树的节点数小于 k-1 ，那么第 k 小的数一定在其右子树上
• 如果等于，那么当前节点就是第 k 小
• 如果大于，那就在左子树上

代码（C++)

中序遍历

class Solution {
public:
    vector<int> nums;
    void dfs(TreeNode *root) {
        if (!root)
            return ;
        dfs(root->left);
        nums.push_back(root->val);
        dfs(root->right);
    }
    int kthSmallest(TreeNode* root, int k) {
        dfs(root);
        return nums[k - 1];
    }
};

优化

这里用的是官方的代码

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

总结

做这题很容易想到给一个数组，求第 K 大/小的值，其中数组里的值可以重复出现。那题也就去重难点，不过如果写过就容易了。这题由于是搜索树，所以无法出现重复值，但是方法二还是很难想到。方法一的递归可以用迭代做，都一样。