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

题目描述

这是 LeetCode 上的 230. 二叉搜索树中第K小的元素 ，难度为 中等。

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

示例 1：

1
2
3
输入：root = [3,1,4,null,2], k = 1

输出：1

示例 2：

1
2
3
输入：root = [5,3,6,2,4,null,null,1], k = 3

输出：3

提示：

  • 树中的节点数为 n 。
  • $1 <= k <= n <= 10^4$
  • $0 <= Node.val <= 10^4$

树的遍历 + 排序

朴素的做法是先对二叉树进行一次完整遍历，将所有节点存入列表中，最后对列表排序后返回目标值。

树的遍历可以使用 DFS 或 BFS。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
    List<Integer> list = new ArrayList<>();
    public int kthSmallest(TreeNode root, int k) {
        dfs(root);
        Collections.sort(list);
        return list.get(k - 1);
    }
    void dfs(TreeNode root) {
        if (root == null) return ;
        list.add(root.val);
        dfs(root.left);
        dfs(root.right);
    }
}

  • 时间复杂度：树的遍历时间复杂度为 $O(n)$；排序的复杂度为 $O(n\log{n})$。整体复杂度为 $O(n\log{n})$
  • 空间复杂度：$O(n)$

树的遍历 + 优先队列（堆）

相比于先直接拿到所有节点再排序的解法一，另外一种做法是使用「优先队列（堆）」来做。

由于我们返回的是第 $k$ 小的数，因此我们可以构建一个容量为 $k$ 的大根堆。

根据大根堆的元素个数和当前节点与堆顶元素的关系来分情况讨论：

  • 大根堆元素不足 $k$ 个：直接将当前节点值放入大根堆；
  • 大根堆元素为 $k$ 个，根据堆顶元素和当前节点值的大小关系进一步分情况讨论：
    • 如果当前节点值元素大于堆顶元素，说明当前节点值不可能在第 $k$ 小的范围内，直接丢弃；
    • 如果当前节点值元素小于堆顶元素，说明当前节点值可能在第 $k$ 小的范围内，先 poll 一个再 add 进去。

树的遍历可以使用 DFS 或 BFS。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution {
    public int kthSmallest(TreeNode root, int k) {
        PriorityQueue<Integer> q = new PriorityQueue<>((a,b)->b-a);
        Deque<TreeNode> d = new ArrayDeque<>();
        d.addLast(root);
        while (!d.isEmpty()) {
            TreeNode node = d.pollFirst();
            if (q.size() < k) {
                q.add(node.val);
            } else if (q.peek() > node.val) {
                q.poll();
                q.add(node.val);
            }
            if (node.left != null) d.addLast(node.left);
            if (node.right != null) d.addLast(node.right);
        }
        return q.peek();
    }
}

  • 时间复杂度：树的遍历时间复杂度为 $O(n)$；使用优先队列（堆）复杂度为 $O(n\log{k})$。整体复杂度为 $O(n\log{k})$
  • 空间复杂度：空间多少取决于 d 和 q 使用的容量，q 最多不超过 $k$ 个元素，复杂度为 $O(k)$，d 最多不超过二叉树的一层，复杂度为 $O(n)$。整体复杂度为 $O(n + k)$

中序遍历

上述两种节点，都没有利用该树为二叉搜索树的特性。

而我们知道，二叉搜索树的中序遍历是有序的，因此我们只需要对二叉搜索树执行中序遍历，并返回第 $k$ 小的值即可。

不熟悉二叉树的中序遍历的同学，可以看看 （题解）783. 二叉搜索树节点最小距离。

中序遍历有「迭代」和「递归」两种写法。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Deque<TreeNode> d = new ArrayDeque<>();
        while (root != null || !d.isEmpty()) {
            while (root != null) {
                d.addLast(root);
                root = root.left;
            }
            root = d.pollLast();
            if (--k == 0) return root.val;
            root = root.right;
        }
        return -1; // never
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
    int k, ans;
    public int kthSmallest(TreeNode root, int _k) {
        k = _k;
        dfs(root);
        return ans;
    }
    void dfs(TreeNode root) {
        if (root == null || k <= 0) return ;
        dfs(root.left);
        if (--k == 0) ans = root.val;
        dfs(root.right);
    }
}
  • 时间复杂度：由于存在对 $k$ 大小的剪枝，因此整个遍历顺序是从小到大进行遍历，搜索到目标值即结束。复杂度为 $O(k)$
  • 空间复杂度：令 $h$ 为树高，复杂度为 $O(h)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.230 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。