【LeetCode-230】二叉搜索树中第K小的元素

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

8.18.1 题目描述

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

8.18.2 方法一：中序遍历

预备知识

二叉搜索树具有如下性质：

• 结点的左子树只包含小于当前结点的数。
• 结点的右子树只包含大于当前结点的数。
• 所有左子树和右子树自身必须也是二叉搜索树。

二叉树的中序遍历即按照访问左子树——根结点——右子树的方式遍历二叉树；在访问其左子树和右子树时，我们也按照同样的方式遍历；直到遍历完整棵树。

思路和算法

因为二叉搜索树和中序遍历的性质，所以二叉搜索树的中序遍历是按照键增加的顺序进行的。于是，我们可以通过中序遍历找到第 k 个最小元素。

「二叉树的中序遍历」可以参考「94. 二叉树的中序遍历的官方题解」，具体地，我们使用迭代方法，这样可以在找到答案后停止，不需要遍历整棵树。

代码

class Solution {
   
    public int kthSmallest(TreeNode root, int k) {
   
        Deque<TreeNode> stack = new ArrayDeque<TreeNode>();
        while (root != null || !stack.isEmpty()) {
   
            while (root != null) {
   
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            --k;
            if (k == 0) {
   
                break;
            }
            root = root.right;
        }
        return root.val;
    }
}

复杂度分析

8.18.3 方法二：记录子树的结点数

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

思路和算法

在方法一中，我们之所以需要中序遍历前 k 个元素，是因为我们不知道子树的结点数量，不得不通过遍历子树的方式来获知。

因此，我们可以记录下以每个结点为根结点的子树的结点数，并在查找第 k 小的值时，使用如下方法搜索：

• 令 node 等于根结点，开始搜索。
• 对当前结点 node 进行如下操作：
  • 如果 node 的左子树的结点数 left 小于 k−1，则第 k 小的元素一定在 node 的右子树中，令 node 等于其的右子结点，k 等于k−left−1，并继续搜索；
  • 如果 node 的左子树的结点数 left 等于 k−1，则第 k 小的元素即为 node ，结束搜索并返回 node 即可；
  • 如果 node 的左子树的结点数 left 大于 k−1，则第 k 小的元素一定在 node 的左子树中，令 node 等于其左子结点，并继续搜索。

在实现中，我们既可以将以每个结点为根结点的子树的结点数存储在结点中，也可以将其记录在哈希表中。

class Solution {
   
    public int kthSmallest(TreeNode root, int k) {
   
        MyBst bst = new MyBst(root);
        return bst.kthSmallest(k);
    }
}

class MyBst {
   
    TreeNode root;
    Map<TreeNode, Integer> nodeNum;

    public MyBst(TreeNode root) {
   
        this.root = root;
        this.nodeNum = new HashMap<TreeNode, Integer>();
        countNodeNum(root);
    }

    // 返回二叉搜索树中第k小的元素
    public int kthSmallest(int k) {
   
        TreeNode node = root;
        while (node != null) {
   
            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;
    }

    // 统计以node为根结点的子树的结点数
    private int countNodeNum(TreeNode node) {
   
        if (node == null) {
   
            return 0;
        }
        nodeNum.put(node, 1 + countNodeNum(node.left) + countNodeNum(node.right));
        return nodeNum.get(node);
    }

    // 获取以node为根结点的子树的结点数
    private int getNodeNum(TreeNode node) {
   
        return nodeNum.getOrDefault(node, 0);
    }
}

复杂度分析

• 时间复杂度：预处理的时间复杂度为 O(N)，其中 N 是树中结点的总数；我们需要遍历树中所有结点来统计以每个结点为根结点的子树的结点数。搜索的时间复杂度为 O(H)，其中 H 是树的高度；当树是平衡树时，时间复杂度取得最小值 O(logN)；当树是线性树时，时间复杂度取得最大值 O(N)。
• 空间复杂度：O(N)，用于存储以每个结点为根结点的子树的结点数。

8.18.4 方法三：平衡二叉搜索树

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

预备知识

方法三需要先掌握 平衡二叉搜索树（AVL树） 的知识。平衡二叉搜索树具有如下性质：

• 平衡二叉搜索树中每个结点的左子树和右子树的高度最多相差 1；
• 平衡二叉搜索树的子树也是平衡二叉搜索树；
• 一棵存有 n 个结点的平衡二叉搜索树的高度是 O(logn)。

思路和算法

我们注意到在方法二中搜索二叉搜索树的时间复杂度为 O(H)，其中 H 是树的高度；当树是平衡树时，时间复杂度取得最小值 O(logN)。因此，我们在记录子树的结点数的基础上，将二叉搜索树转换为平衡二叉搜索树，并在插入和删除操作中维护它的平衡状态。

其中，将二叉搜索树转换为平衡二叉搜索树，可以参考「1382. 将二叉搜索树变平衡的官方题解」。在插入和删除操作中维护平衡状态相对复杂，读者可以阅读下面的代码和注释，理解如何通过旋转和重组实现它。

class Solution {
   
    public int kthSmallest(