Top-K 算法的实现：从堆到快速选择
七崽爱吃小饼干2026/10/08阅读 0
面试里"求第 K 大""找 Top K"出现的频率极高，但很多人第一反应就是一句 sort 了事。排序当然能解，可当数据是百万级、甚至是一个流、根本装不下内存时，O(N log N) 的全排序又慢又浪费。
Top-K 的本质是：我只要最大的 K 个，不需要把 N 个全排好。这篇讲清几种主流实现、各自的复杂度，以及怎么选。
一、问题定义
输入：一个数组 nums 和整数 K；输出：最大的 K 个元素（或第 K 大）。
最直白的解法是先排序再取前 K 个：
js
function topK(nums, k) {
return [...nums].sort((a, b) => b - a).slice(0, k)
}
- 时间：
O(N log N)，空间O(N)（拷贝）。 - 优点：三行搞定，
N小或本来就要完整排序时，这是最优解，别过度设计。 - 缺点：为了 K 个元素把所有元素都排了，
K很小、N很大时纯属浪费。
更好的思路有三类：堆、快选、计数。
二、三种思路一览
| 方法 | 平均时间 | 空间 | 适合 |
|---|---|---|---|
| 全排序 | O(N log N) | O(N) | N 小、需要完整顺序 |
| 小顶堆 | O(N log K) | O(K) | N 大、K 小、流式数据 |
| 快速选择 | O(N)（最坏 O(N²)） | O(1) | 一次性、可随机访问的数组 |
| 计数/桶 | O(N + 值域) | O(值域) | 值域小的整数 |
下面逐个实现。
三、小顶堆：最通用，适合大数据和流
要维护"当前最大的 K 个"，用一个容量为 K 的小顶堆：堆顶是这 K 个里最小的。
遍历每个元素：
- 堆没满 → 直接入堆。
- 堆满了 → 若当前元素比堆顶大，就把堆顶弹掉、把当前元素入堆。
这样堆里永远只留最大的 K 个，堆顶是它们的"门槛"。复杂度 O(N log K)，空间 O(K)，数据流式和内存受限场景的首选。
先实现一个最小堆（JS 没有内置）：
js
class MinHeap {
constructor() { this.a = [] }
size() { return this.a.length }
peek() { return this.a[0] }
push(v) {
const a = this.a
a.push(v)
let i = a.length - 1
while (i > 0) {
const p = (i - 1) >> 1
if (a[p] <= a[i]) break
;[a[p], a[i]] = [a[i], a[p]]
i = p
}
}
pop() {
const a = this.a
const top = a[0]
const last = a.pop()
if (a.length) {
a[0] = last
let i = 0
for (;;) {
const l = i * 2 + 1
const r = l + 1
let m = i
if (l < a.length && a[l] < a[m]) m = l
if (r < a.length && a[r] < a[m]) m = r
if (m === i) break
;[a[i], a[m]] = [a[m], a[i]]
i = m
}
}
return top
}
}
function topKByHeap(nums, k) {
if (k <= 0) return []
const heap = new MinHeap()
for (const n of nums) {
if (heap.size() < k) {
heap.push(n)
} else if (n > heap.peek()) {
heap.pop()
heap.push(n)
}
}
// 堆内是最大的 K 个，弹出即降序
const res = []
while (heap.size()) res.push(heap.pop())
return res.reverse()
}
console.log(topKByHeap([3, 1, 4, 1, 5, 9, 2, 6], 3)) // [9, 6, 5]
要点：
- 找 Top K 大用小顶堆，找 Top K 小用大顶堆——堆顶始终是"要被淘汰的那一个"。
- 结果顺序：直接弹堆得到的是升序，想要降序再
reverse。 - 想找第 K 大，堆顶就是答案，连弹都不用。
四、快速选择：平均 O(N)，一次性求第 K 大
快速选择（Quickselect）借用快排的 partition 思想：随便选个 pivot，把数组分成"小于 pivot"和"大于 pivot"两半，然后只往目标 K 所在的那一半继续分区，另一半直接丢弃。
和快排的区别就在这：快排两边都要递归，快选只递归一边，所以平均复杂度从 O(N log N) 降到 O(N)。
js
function quickSelect(nums, k) {
// 返回第 k 大（k 从 1 开始）
const a = [...nums]
const target = a.length - k // 升序下的目标下标
let lo = 0
let hi = a.length - 1
while (lo < hi) {
const p = partition(a, lo, hi)
if (p === target) break
else if (p < target) lo = p + 1
else hi = p - 1
}
return a[target]
}
function partition(a, lo, hi) {
const pivot = a[hi]
let i = lo
for (let j = lo; j < hi; j++) {
if (a[j] < pivot) {
;[a[i], a[j]] = [a[j], a[i]]
i++
}
}
;[a[i], a[hi]] = [a[hi], a[i]]
return i
}
console.log(quickSelect([3, 1, 4, 1, 5, 9, 2, 6], 3)) // 6，第三大
坑点：
- 最坏 O(N²)：每次都取到极值当 pivot（比如有序数组）。工程上用随机 pivot 或 三数取中 规避。
- 返回第 K 大后，
a.slice(target)再排一下就是 Top K。 - 它需要能随机访问整个数组，流式数据用不了。
五、计数/桶：值域小的时候 O(N)
如果数据是值域很小的整数（比如成绩 0100、年龄、评分 15），直接用数组下标计数，一次遍历搞定：
js
function topKByCount(nums, k, maxValue = 100) {
const count = new Array(maxValue + 1).fill(0)
for (const n of nums) count[n]++
const res = []
for (let v = maxValue; v >= 0 && res.length < k; v--) {
for (let c = 0; c < count[v] && res.length < k; c++) res.push(v)
}
return res
}
时间 O(N + 值域)，比堆还快，但只适用于值域可控的整数。值域巨大时，可退化为分桶 + 桶内再选。
六、海量/流式数据怎么办
当数据超过单机内存，或是一条永远不结束的流：
- 流式：用堆（
O(K)内存）或 Count-Min Sketch + 堆 估高频项（Redis 的热 key、热搜榜思路）。 - 分布式：Map 阶段每个分片各求自己的 Top K，Reduce 阶段合并这些候选再求全局 Top K。因为全局 Top K 一定包含在各分片 Top K 的并集里——这一步的正确性很关键。
- 外部排序：数据落盘时，分块排序后多路归并。
七、怎么选
| 需求 | 选择 |
|---|---|
| N 小、要完整顺序 | 直接 sort |
| N 大、K 小、通用 | 小顶堆 |
| 流式 / 内存受限 | 小顶堆 |
| 一次性求第 K 大 | 快速选择 |
| 只需第 K 大且数据可随机访问 | 快速选择 or 堆 |
| 值域小的整数 | 计数/桶 |
| 海量分布式 | 分片 Top K + 归并 |
一个常被忽略的优化：当 K 接近 N 时，堆的 O(N log K) 反而不划算，直接排序更省事。反之 K 远小于 N 时堆优势才明显。
小结
- Top-K 不必全排序：只要最大的 K 个，别为 K 个元素排 N 个。
- 小顶堆求 Top K 大，
O(N log K)、O(K)空间，是通用解和流式数据的首选；求第 K 大时堆顶即答案。 - 快速选择平均
O(N)，适合一次性、可随机访问的数组，最坏O(N²)，要随机化 pivot。 - 计数/桶在值域小的整数上能做到
O(N)。 - 海量/流式用堆 + 分片归并（或 Count-Min Sketch）解决。
- 选型口诀：能全排就全排，要省就上堆，求一个数用快选，整数小值域就计数。