LeetCode 剑指 Offer 40. 最小的k个数
题目描述
题意分析
从数组
arr中选出最小的k个数并返回，结果可以按任意顺序排列。每个数按在原数组中的出现次数参与选择，相同值出现多次时，也可能在答案中保留多次，不能先去重。题目保证
0 <= k <= arr.length。k = 0时返回空结果，k等于数组长度时返回全部元素。需要的是最小的这批数，不要求把整个数组排好序，也不要求只返回第k小的单个值。
解法：大小为 k 的最大堆
核心思路
[!blue]
扫描数组时，维护目前见过的最小
k个候选。需要随时淘汰的是候选中最大的那个，所以使用最大堆，让它位于堆顶。堆内不足k个时直接加入；堆满后才需要决定新数是否能进入答案。设堆顶为
largest。若新数x >= largest，堆中的k个数已经都不大于x，没有必要将它加入；若x < largest，应删除当前最大的候选并加入x。这样每处理一个新数，堆仍保存已处理部分中最小的min(已处理数量, k)次出现，最终就是整个数组的答案。相等时不替换，只是用已有的一次出现代表同值的新出现，不会减少候选数量；未满时遇到重复值仍然正常入堆。因此这个过程保留的是按次数计数的元素，而不是不同数值的集合。
Java 最后逐次弹出最大值，得到降序结果；Go 直接返回堆中的存储，只保证父子之间的堆序。两者都满足任意顺序的要求，无需为了输出再排序。这种方法只读取输入，额外维护
k个候选，也适合数据逐个到达的场景。
解题步骤
k == 0时直接返回空结果，避免查询空堆顶。- 创建最大堆，逐个读取数组元素。
- 堆未满时加入当前数；已满时，只在当前数小于堆顶时替换堆顶。
- 扫描结束后返回堆中的全部
k个数。
代码实现
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;
}
}
import "container/heap"
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)
}
复杂度分析
- 时间复杂度：
k > 0时为 $O(n \log(k + 1))$，堆最多含k个元素。Java 最后的出堆共 $O(k \log(k + 1))$，已包含在总上界内；Go 直接返回堆存储。k = 0时为 $O(1)$。- 空间复杂度：$O(k)$，用于堆和输出，不修改输入数组。
关键点总结
[!green]
- 求最小的一批数，用最大堆暴露最该被淘汰的候选。
- 堆未满时保留所有出现，堆满后才按新值是否更小决定替换。
- 堆内只需保证候选集合正确，不需要让它们完全有序。
解法二：随机三路快速选择
核心思路
[!blue]
如果允许重排输入，可以只把最小的
k个数分到前面，不必让前后两部分各自有序。目标边界是按升序排列时的下标k - 1，利用快速排序的划分操作即可逐步定位它。在当前区间随机选择一个基准值，用三路划分把元素分成小于、等于、大于基准的三段。维护
[left, less)小于基准、[less, index)等于基准、[index, greater]尚未处理、(greater, right]大于基准。遇到小值时与less交换并推进两个下标；遇到大值时与greater交换，只收缩右边界；遇到相等值时只推进扫描下标。划分结束后，
[less, greater]就是相等段。若k - 1 < less，边界在小于段，只需继续处理左边；若k - 1 > greater，边界在大于段，只需继续处理右边；否则，前k个位置已经包含全部更小值和足够多的相等值，可以结束。每次都只进入包含目标边界的一侧，另一侧无需继续排序。随机基准降低持续出现极端不平衡划分的概率，三路划分则一次处理全部相等元素。最终返回数组前
k项；它们属于正确的最小集合，但内部顺序不限。Java 复制前k项作为返回数组，Go 直接返回相应切片，两种实现都会重排原数组。
解题步骤
k = 0时返回空结果，否则令目标下标为k - 1。- 在当前区间随机取基准值，完成一次三路划分。
- 根据目标下标与相等段的位置关系，收缩到左侧或右侧；目标落在相等段时停止。
- 返回数组前
k个元素。
代码实现
class Solution {
public int[] getLeastNumbers(int[] arr, int k) {
if (k == 0) {
return new int[0];
}
int left = 0;
int right = arr.length - 1;
while (left < right) {
int pivot = arr[left + (int) (Math.random() * (right - left + 1))];
int less = left;
int index = left;
int greater = right;
while (index <= greater) {
if (arr[index] < pivot) {
swap(arr, index++, less++);
} else if (arr[index] > pivot) {
swap(arr, index, greater--);
} else {
index++;
}
}
if (k - 1 < less) {
right = less - 1;
} else if (k - 1 > greater) {
left = greater + 1;
} else {
break;
}
}
return Arrays.copyOf(arr, k);
}
private void swap(int[] arr, int i, int j) {
int value = arr[i];
arr[i] = arr[j];
arr[j] = value;
}
}
import "math/rand"
func getLeastNumbers(arr []int, k int) []int {
if k == 0 {
return arr[:0]
}
left, right := 0, len(arr)-1
for left < right {
pivot := arr[left+rand.Intn(right-left+1)]
less, index, greater := left, left, right
for index <= greater {
if arr[index] < pivot {
arr[index], arr[less] = arr[less], arr[index]
index++
less++
} else if arr[index] > pivot {
arr[index], arr[greater] = arr[greater], arr[index]
greater--
} else {
index++
}
}
if k-1 < less {
right = less - 1
} else if k-1 > greater {
left = greater + 1
} else {
break
}
}
return arr[:k]
}
复杂度分析
- 时间复杂度：期望 $O(n)$，每次划分只继续处理一侧，随机基准使期望处理量线性；最坏划分仍可能达到 $O(n^2)$。Java 复制结果另需 $O(k)$，不改变期望上界。
- 空间复杂度：$O(1)$，不计返回结果。原地划分且使用循环，不需要递归栈；Java 返回数组占 $O(k)$，Go 返回切片共享输入存储。
关键点总结
[!green]
- 只定位第
k小的边界，不对两侧完成全排序。- 相等段一次处理重复值，目标落入其中即可结束。
- 以重排输入换取常数辅助空间；随机化改善期望复杂度，不提供最坏线性保证。
易错点总结
[!yellow]
- 堆方向写反，堆顶会变成最应该保留的最小数，替换时就删错候选。
- 堆满后无条件替换，会让更大的新数挤掉本应保留的元素。
- 忘记处理
k = 0，会访问空堆顶；快速选择时也会产生无效的目标下标-1。- 对输入去重会丢失重复元素的次数，无法保证返回正确的
k个数。- 快速选择时把
k当成目标下标，会偏移一位；边界应围绕k - 1判断。- 三路划分中把大元素换到右边后，换回的元素仍未分类，当前扫描下标不能立刻增加。
- 快速选择会重排输入，Go 返回的切片还与输入共享存储；需要保留输入时使用上面的堆解法。
相似题目
| 题目 | 难度 | 关联与区别 |
|---|---|---|
| 703. 数据流中的第 K 大元素 | 简单 | 原题持续加入数据并查询第k大，本题只做一次选择，可用快速选择或堆。 |
| 347. 前 K 个高频元素 | 中等 | 同样挑选TopK，原题排名依据为元素频次，本题依据数值大小。 |
| 补充题 208. 最大的 K 个数 | 简单 | 都用容量为 k 的堆维护候选；本题用大顶堆保留最小值，补充题用小顶堆保留最大值。 |