16. 找到堆的前 k 个元素

0 - 算法, 0.0 - 算法模版 本文总阅读量次

正文

“找到第 K 大的元素”或“找到前 K 个高频元素”是堆（Heap/PriorityQueue）最经典的应用场景。利用堆的特性，可以将原本需要全排序的 $O(N \log N)$ 复杂度优化到 $O(N \log K)$。

算法模板 (找到第 K 大的元素)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import java.util.PriorityQueue

fun findKthLargest(nums: IntArray, k: Int): Int {
    // 技巧：求第 K 大用最小堆；求第 K 小用最大堆
    // 创建最小堆 (PriorityQueue 默认即为最小堆)
    val minHeap = PriorityQueue<Int>()

    // 将数组中的元素依次加入最小堆
    for (num in nums) {
        minHeap.offer(num)

        // 如果最小堆的大小超过 k，移除堆顶元素
        // 这样堆中始终保留的是目前遍历过的最大的 k 个元素
        if (minHeap.size > k) {
            minHeap.poll()
        }
    }

    // 最小堆的堆顶就是这 k 个最大值中最小的一个，即第 k 大的元素
    return minHeap.peek()
}

使用说明

  1. 核心逻辑：

    • 求前 K 大：维护一个大小为 $K$ 的最小堆。当堆满时，新元素若大于堆顶，则弹出堆顶并插入新元素。最终堆中留下的就是最大的 $K$ 个数，堆顶是第 $K$ 大。
    • 求前 K 小：维护一个大小为 $K$ 的最大堆（PriorityQueue(Collections.reverseOrder())）。堆顶是这 $K$ 个最小值中最大的一个，即第 $K$ 小。

  2. 核心优势：

    • 空间效率：不需要加载全部数据，只需维持 $K$ 个元素的堆。
    • 时间效率：$O(N \log K)$。当 $K \ll N$ 时，性能接近 $O(N)$。

  3. 常见变形：

    • 前 K 个高频元素：配合哈希表记录频率，堆中存储 Map.Entry 并根据频率排序。
    • 合并 K 个有序链表：利用最小堆维护 $K$ 个链表的当前头节点。
    • 数据流中的第 K 大：保持堆的大小为 $K$，持续 offer 和 poll。

  4. 注意点：

    • 在 Kotlin/Java 中，PriorityQueue 默认是最小堆（队首是最小值）。
    • 如果需要最大堆，需传入比较器：PriorityQueue<Int> { a, b -> b - a }。
    • 对于海量数据处理，堆是内存占用最友好的方案。
0 - 算法, 0.0 - 算法模版