面试：Top K问题

什么是 Top K？想象一下这些场景：微博的热搜榜单、电商平台的销量前十商品、游戏中全服战力排名前一百的玩家，或者从几十亿条搜索日志中找出搜索频率最高的 10 个词。

在一个海量的数据集中，找出最大（或最小、最常出现）的前 $K$ 个元素，这就是大名鼎鼎的 Top K 问题。

为什么不用普通排序？

面对这个问题，很多同学的第一反应是：直接把整个数组按降序排好，然后取前 $K$ 个不就行了？

想法没错，但如果数据量极其庞大（比如有 10 亿个数），而 $K$ 只有 10，把 10 亿个数全部排序是非常浪费计算资源和内存的。我们需要一种更高效、甚至能够处理“动态流数据”的解决方案。

目前主流的解法有两种：堆（Priority Queue）和快速选择（Quick Select）。今天我们重点讲解在工程中最常用、最适合处理海量流数据的堆排序法。

核心思想：小顶堆（Min-Heap）的妙用

核心逻辑如下：

1、我们维护一个容量只有 $K$ 的小顶堆。小顶堆的特点是：堆顶元素是整个堆里最小的。

2、遍历数据集，当堆里的元素不足 $K$ 个时，直接把元素塞进去。

3、当堆满了 $K$ 个元素后，新来的元素必须和“堆顶元素（当前 K 个数中最小的那个）”进行比较。如果新元素比堆顶还小，直接淘汰；如果新元素比堆顶大，说明堆顶那个家伙不配留在 Top K 里，我们把堆顶踢走，让新元素进堆。

4、遍历完所有数据后，留在堆里的这 $K$ 个元素，就是我们要找的 Top K！

代码实现

在 Java 中，完美契合“堆”这一数据结构的就是优先队列 PriorityQueue。它默认就是一个小顶堆，这让我们的代码编写变得异常简洁。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import java.util.PriorityQueue;
import java.util.Arrays;

public class TopKDemo {

    public static void main(String[] args) {
        int[] data = {3, 1, 5, 12, 2, 11, 7, 6, 4, 15};
        int k = 3;
        
        System.out.println("原始数据：" + Arrays.toString(data));
        int[] topK = findTopKLargest(data, k);
        System.out.println("最大的 " + k + " 个数是：" + Arrays.toString(topK));
    }

    /**
     * 基于小顶堆寻找最大的 K 个元素
     * @param nums 原始数组
     * @param k    Top K 的 K 值
     * @return     包含最大的 K 个元素的数组
     */
    public static int[] findTopKLargest(int[] nums, int k) {
        // 参数校验
        if (nums == null || nums.length == 0 || k <= 0) {
            return new int[0];
        }

        // Java 的 PriorityQueue 默认就是小顶堆
        PriorityQueue<Integer> minHeap = new PriorityQueue<>(k);

        // 遍历所有数据
        for (int num : nums) {
            if (minHeap.size() < k) {
                // 如果堆还没满，直接加入
                minHeap.offer(num);
            } else if (num > minHeap.peek()) {
                // 如果堆满了，且当前元素大于堆顶（堆中最小元素）
                // 则将堆顶踢出，将新元素加入
                minHeap.poll();
                minHeap.offer(num);
            }
        }

        // 将堆中的结果转存到数组中并返回
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            // 注意：此时取出的顺序是从小到大
            result[i] = minHeap.poll();
        }
        
        return result;
    }
}