LeetCode 剑指 Offer 40. 最小的k个数

约 6 分钟
文章信息 19 个分类 7 个标签
分类
LeetCode CodeTop 简单 已解决 已掌握 20-50 字节跳动 微软 美团 阿里巴巴 快手 腾讯 猿辅导 百度 滴滴 小米 商汤 剑指 Offer 剑指 Offer（第 2 版）
标签
数组 排序 堆（优先队列） 分治 Top N 问题 快速选择 大顶堆
目录

题目描述

✅ 剑指 Offer 40. 最小的k个数

image-20250420081543690

image-20241107211133333

题意分析

给定一个整数数组 arr 和一个整数 k，返回数组中最小的 k 个数组成的数组。

最关键的一句话藏在题面里：返回的这 k 个数可以按任意顺序排列。这句话把问题的性质彻底改变了——它要的是一个「集合」而不是一个「有序序列」，也就是说我们根本不需要知道这 k 个数彼此之间谁大谁小，只需要把它们和剩下的 n - k 个数分开。凡是把「取最小的 k 个」条件反射成「排序」的，都是在做题目没要求的额外功课。

另一个要留意的信号是：题目只关心这 k 个数是谁，不关心第 k 小具体是多少，也不要求保持原数组顺序，因此可以自由地打乱、交换、原地修改输入数组。

边界上必须点名的是 k 可以取 0，此时应当返回空数组，任何直接访问 arr[0] 或建堆后取堆顶的写法都会踩空；另一端 k 可以等于 arr.length，这时整个数组都是答案。数组元素可能重复，重复元素各算各的，例如 [1, 1, 2] 取 k = 2 的答案是两个 1。

解法：大小为 k 的最大堆

核心思路

问题关键：题目只要最小的 k 个数，且不要求答案有序。全量排序会额外得到所有元素的次序，花费 $O(n \log n)$；真正需要维护的只有 k 个候选及其边界。

为什么选最大堆：候选集合中最可能被淘汰的是最大值，因此用大小为 k 的最大堆，让堆顶始终表示当前入选门槛。新数小于堆顶时替换堆顶，否则它不可能进入前 k 小。这里方向容易记反：保留最小的 k 个，要让候选中的最大值位于堆顶。

不变量：处理完前 i 个数后，堆中恰好保存其中最小的 min(i, k) 个数，堆满时堆顶是候选中的最大值。

正确性：堆未满时直接加入；堆满后，若 x >= top，当前已有 k 个数不大于 x，丢弃 x 不影响答案；若 x < top，用 x 替换候选中最大的 top，新的候选仍是前 i + 1 个数中最小的 k 个。因此遍历结束时堆中就是答案。

快速选择平均可做到 $O(n)$，但会修改数组且最坏为 $O(n^2)$；最大堆上界稳定，还能处理数据流，是本题更稳妥的面试主解。

解题步骤

  1. k == 0 时直接返回空数组，避免访问空堆的堆顶。
  2. 创建最大堆，依次扫描数组；堆不足 k 个时直接入堆。
  3. 堆已满时，仅当当前值小于堆顶，才弹出堆顶并加入当前值。
  4. 扫描结束后取出堆中全部元素；题目允许任意顺序，无需再排序。

口述样例：[3, 2, 1]，k = 2。先将 3、2 入堆，堆顶为 3；读到 1 时用它替换 3，最终堆中是 {1, 2}。

代码实现

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] res = new int[k];
        if (k == 0) {
            return res;
        }

        PriorityQueue<Integer> heap = new PriorityQueue<>(
                (a, b) -> Integer.compare(b, a));
        for (int x : arr) {
            if (heap.size() < k) {
                heap.offer(x);
            } else if (x < heap.peek()) {
                heap.poll();
                heap.offer(x);
            }
        }

        for (int i = 0; i < k; i++) {
            res[i] = heap.poll();
        }
        return res;
    }
}

type maxHeap []int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any)        { *h = append(*h, x.(int)) }

func (h *maxHeap) Pop() any {
    old := *h
    n := len(old)
    v := old[n-1]
    *h = old[:n-1]
    return v
}

func getLeastNumbers(arr []int, k int) []int {
    if k == 0 {
        return []int{}
    }

    h := &maxHeap{}
    for _, x := range arr {
        if h.Len() < k {
            heap.Push(h, x)
        } else if x < (*h)[0] {
            heap.Pop(h)
            heap.Push(h, x)
        }
    }
    return []int(*h)
}

复杂度分析

  • 时间复杂度：$O(n \log k)$。扫描 n 个元素，堆的大小始终不超过 k，每次调整为 $O(\log k)$。
  • 空间复杂度：$O(k)$。最大堆只保存 k 个候选；返回数组不计入额外空间。

关键点总结

  • 堆顶应放「当前最该淘汰的候选」：求最小的 k 个用最大堆，求最大的 k 个用最小堆。
  • 堆必须限制为 k 个元素，复杂度中的对数项才是 $\log k$ 而不是 $\log n$。
  • 题目不要求顺序，输出堆内容即可；额外排序只会增加 $O(k \log k)$。
  • 面试追问快速选择时，要能说明它平均更快，但会改动输入、最坏复杂度不稳定，也不适合流式数据。

易错点总结

  • 堆方向写反：[3,2,1]，k=2 若用最小堆按同样逻辑淘汰，会错误保留 {2,3}。
  • 堆满后无条件替换：[1,2,3]，k=2 会把本应保留的 2 换成 3；必须先比较 x < heap.peek()。
  • 漏掉 k == 0：最终弹出空堆会发生空指针或越界。
  • Java 比较器写成 b - a：范围扩大后可能整数溢出，使用 Integer.compare(b, a) 更稳妥。
  • Go 的 Less 必须使用 > 才是最大堆，且修改切片长度的 Push、Pop 必须使用指针接收者。

相似题目

题目 难度 考察点
215. 数组中的第K个最大元素 中等 堆方向相反：用最小堆求第 K 大的单个值
面试题 17.14. 最小K个数 中等 本题的同题异构，可直接复用最大堆写法
347. 前 K 个高频元素 中等 先哈希统计频次，再按频次而非数值本身建堆
692. 前K个高频单词 中等 频次相同时需按字典序，比较器要写成复合规则
973. 最接近原点的 K 个点 中等 比较键是平方距离，展示如何为对象自定义权重
703. 数据流中的第 K 大元素 简单 数据流场景，凸显定容堆相对排序的核心优势
LCR 076. 数组中的第 K 个最大元素 中等 215 的 LCR 版本，可练手快速选择的分治划分
LCR 060. 前 K 个高频元素 中等 347 的 LCR 版本，适合对比桶排序与堆两种解法