从海量数字中寻找最大/小的 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 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 list = new ArrayList<>();
Iterator 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 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 list = new ArrayList<>();
Iterator 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));
}