欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

查找 TopK 问题

程序员文章站 2024-03-15 20:28:48
...

​从海量数字中寻找最大/小的 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));
}

更多文章欢迎扫码关注公众号 [JAVA觅音阁] 查看
查找 TopK 问题