算法题分享 | 二叉搜索树中第 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 小的值，最简单的做法是直接使用中序遍历，第 K 个遍历到的数即为第 K 小的数。

这种方法虽然实现起来最为简单，但是由于每次查找都要中序遍历，所以对于要频繁查找第 K 小的值，性能是不够理想的，我们可以在此基础上继续优化。

上面的方法中之所以每次都要遍历 K 个元素，是因为我们不知道每个子树的节点数，如果知道以每个结点为根的子树的结点数，那么就可以结合结点数信息在更小的时间复杂度内找到第 K 小元素。例如，对于结点 node, 若其左子树结点数大于 K-1，即说明目标结点在左子树中，进而往左子树找即可；若其左子树结点数等于 k-1，即表明 node 就是目标节点，返回其值即可；而左子树结点数小于 K-1，就说明目标结点在右子树当中，需要注意的是，在往右子树查找前，需要先更新 k 值为 k - (左子树结点数 + 1(当前结点))。
使用这种方法需要在初始化时统计并记录各结点对应子树的总结点数目。
本文给出的代码即为这种方法的实现。

拓展 对于需要经常修改（插入/删除操作）二叉搜索树的场景，由于普通的二叉搜索树是无法保证平衡的，所以在极端的情况下（线性树）是会导致查询时间复杂度恶化的，所以最好是在上面所说的记录子树结点数的基础上，使用平衡二叉搜索树进行处理，在初始化时根据原树构建出平衡二叉搜索树，而对于子树结点数，即可以直接存储在各结点中。
通过实现平衡二叉查找树，可以使得查找时间复杂度稳定在 O(logn)。

代码

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public int kthSmallest(TreeNode root, int k) {
        MyBst myBst = new MyBst(root);
        return myBst.getTargetVal(k);
    }
}


class MyBst {
    
    TreeNode root;

    Map<TreeNode, Integer> countMap;

    public MyBst(TreeNode root) {
        this.root = root;
        this.countMap = new HashMap<>();
        getNodeCount(root);
    }

    // 获取第 K 小节点值
    public int getTargetVal(int k) {
        TreeNode node = root;
        while (node != null) {
            int count = countMap.getOrDefault(node.left, 0);
            if (count == k - 1) {
                break;
            } else if (count < k - 1) {
                k -= count + 1;
                node = node.right;
            } else {
                node = node.left;
            }
        }

        return node.val;
    }

    private int getNodeCount(TreeNode root) {
        if (root == null) {
            return 0;
        }

        countMap.put(root, getNodeCount(root.left) + getNodeCount(root.right) + 1);

        return countMap.get(root);
    }
}

复杂度分析

• 时间复杂度：预处理需要 O(n) 的时间复杂度，n 为树的结点数。后续搜索的时间复杂度为 O(H)，H 为树的高度。

• 空间复杂度：O(n)
  需要存储各子树的结点数

优质项目推荐

推荐一个可用于练手、毕业设计参考、增加简历亮点的项目。

lemon-puls/txing-oj-backend: Txing 在线编程学习平台，集在线做题、编程竞赛、即时通讯、文章创作、视频教程、技术论坛为一体

公众号

有兴趣可以关注公众号一起学习更多的干货哈！