力扣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^4
• 0 <= Node.val <= 10^4

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

 解法一：中序遍历

/**
 * 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) {
        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;
    }
};

这段代码实现了在二叉搜索树中找到第K小的元素。代码的解析如下：

1. 首先定义了一个栈 stack，用于存放节点。
2. 进入循环，循环条件为 root 不为空或者栈 stack 不为空。
3. 在循环中的第一个内部循环中，将当前节点 root 及其左子节点依次入栈，并更新当前节点为其左子节点。这样可以实现中序遍历的逻辑，即先遍历左子树。
4. 内部循环结束后，表示当前节点及其左子树已经遍历完毕，需要将当前节点出栈，并将 k 减一。这样可以保证找到第K小的元素时，k 的值为0。
5. 判断 k 的值是否为0，如果是，则表示已经找到第K小的元素，退出循环。
6. 否则，需要继续遍历当前节点的右子节点，并将当前节点更新为其右子节点。
7. 循环结束后，返回当前节点的值，即为第K小的元素。

代码的时间复杂度平均情况下为O(logN)，最坏情况下为O(N)，其中N为二叉搜索树的节点数。空间复杂度平均情况下为O(logN)，最坏情况下为O(N)，其中N为二叉搜索树的节点数。

解法二：记录子树结点数

/**
 * 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 mybst {
public:
    mybst(TreeNode* root) {
        this->root = root;
        cuuntnodenum(root);
    }

    int kthSmallest(int k) {
        TreeNode* node = root;
        while (node) {
            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;

    int cuuntnodenum(TreeNode* node) {
        if (node == nullptr)
            return 0;
        nodenum[node] =
            1 + cuuntnodenum(node->left) + cuuntnodenum(node->right);
        return nodenum[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 m(root);
        return m.kthSmallest(k);      
    }
};

 

这段代码实现了在二叉搜索树中找到第K小的元素。代码的思路如下：

1. 首先定义了一个辅助类 mybst，用于处理二叉搜索树的节点数量。
2. mybst 类的构造函数接受一个二叉搜索树的根节点，将其保存在 root 成员变量中，并调用 cuuntnodenum 方法计算每个节点的子节点数量。
3. kthSmallest 方法接受一个整数 k，表示要找到第K小的元素。
4. 在 kthSmallest 方法中，使用一个指针 node 来遍历二叉搜索树。通过循环，不断更新指针 node，直到找到第K小的元素。
5. 在循环中，首先计算当前节点 node 的左子树的节点数量，并保存在变量 left 中。
6. 判断 left 的值与 k-1 的关系：
   • 如果 left 小于 k-1，说明第K小的元素在当前节点的右子树中，将指针 node 更新为当前节点的右子节点，并更新 k 的值，减去 left+1。
   • 如果 left 等于 k-1，说明第K小的元素就是当前节点，退出循环。
   • 如果 left 大于 k-1，说明第K小的元素在当前节点的左子树中，将指针 node 更新为当前节点的左子节点。
7. 循环结束后，返回当前节点 node 的值，即为第K小的元素。

最后，主类 Solution 中的 kthSmallest 方法使用辅助类 mybst 来找到第K小的元素，并返回结果。

这段代码的时间复杂度为O(logN)，其中N为二叉搜索树的节点数。空间复杂度为O(N)，其中N为二叉搜索树的节点数，主要用于保存每个节点的子节点数量。