查找 TopK 问题

从海量数字中寻找最大/小的 k 个，这类问题我们称为 TopK 问题。

通常使用数据结构-最大/小堆来解决

求前 k 大，用最小堆，即堆顶元素为堆中最小值。

求前 k 小，用最大堆，即堆顶元素为堆中最大值。

如前k大的值, 传入列表 list=[12, 39, 3, 72, 56, 81, 15, 9, 103] 和 k=3, 输出 [103, 81, 72]。
如前k小的值, 传入列表 list=[12, 39, 3, 72, 56, 81, 15, 9, 103] 和 k=3, 输出 [3, 9, 12]。

前K大值

思路：

迭代列表元素：

1.先放入元素前 k 个建立一个最小堆；

2.当前元素x若大于堆顶元素：移除堆顶元素并入队x；

3.最后获取最小堆中的值，即为 topK4Max。

JAVA参考代码

public static int[] topK4Max(int[] nums, int k) {

    // 优先队列、从小到大 (默认也是自然排序)
    Queue<Integer> minHeap = new PriorityQueue<>(Comparator.naturalOrder());
    // 建立最小堆
    for (int i = 0, len = nums.length; i < len; i++) {
        if (minHeap.size() < k) {
            minHeap.offer(nums[i]); // 入队
            continue;
        }
        if (nums[i] > minHeap.peek()) { // 取队首元素比较，不移除
            minHeap.poll(); // 取队首元素，并移除
            minHeap.offer(nums[i]);
        }
    }
    // 遍历队列
    List<Integer> list = new ArrayList<>();
    Iterator<Integer> iterator = minHeap.iterator();
    while (iterator.hasNext()){
        list.add(iterator.next());
    }
    // 从大到小排序
    Collections.sort(list, (o1, o2) -> o1 < o2 ? 1 : o1 == o2 ? 0 : -1);
    // 转换成数组
    return list.stream().mapToInt(Integer::valueOf).toArray();

}

前K小值

思路：

迭代列表元素：

1.先放入元素前 k 个建立一个最大堆；

2.当前元素x若小于堆顶元素：移除堆顶元素并入队x；

3.最后获取最大堆中的值，即为 topK4Min。

JAVA参考代码

public static int[] topK4Min(int[] nums, int k) {

    // 优先队列、从大到小排序
    Queue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
    // 建立最大堆
    for (int i = 0, len = nums.length; i < len; i++) {
        if (maxHeap.size() < k) {
            maxHeap.offer(nums[i]); // 入队
            continue;
        }
        if (nums[i] < maxHeap.peek()) { // 取队首元素比较，不移除
            maxHeap.poll(); // 取队首元素，并移除
            maxHeap.offer(nums[i]);
        }
    }
    // 遍历队列
    List<Integer> list = new ArrayList<>();
    Iterator<Integer> iterator = maxHeap.iterator();
    while (iterator.hasNext()){
        list.add(iterator.next());
    }
    // 从小到大排序
    Collections.sort(list);
    // 转换成数组
    return list.stream().mapToInt(Integer::valueOf).toArray();
    
}

测试代码

public static void main(String[] args) {
    int[] nums = new int[]{12, 39, 3, 72, 56, 81, 15, 9, 103};

    int[] maxK = topK4Max(nums, 3);
    System.out.println(JSON.toJSONString(maxK));

    int[] minK = topK4Min(nums, 3);
    System.out.println(JSON.toJSONString(minK));
}