在Java中使用堆排序求解TopK问题

在Java中使用堆排序求解TopK问题

1. 问题描述

给定一个很大的数组，长度100w，求第k大的数是多少？

这个问题是一个很经典的问题，如果采用传统方式，即现排序，然后找到第k个数，对于数据量很大的时候，是非常耗时的。对于这个问题，最好的解决方案是使用堆排序。

2. 暴力解法

最容易想到的办法就是暴力方法，直接将这个数组进行排序，然后直接输出第k-1位置上的元素即可。这种方法对于数据量小的时候，是最简单并且使用的方法，但是设想一下，当数据量高达千万一别的时候，对这些数据排序所话费的时间是很恐怖的。

3. 使用堆解决

我们在学习数据结构的时候学习过，堆在求解第TopK问题的时候是非常高效的，因为他不需要对所有的数据都进行排序，具体的方法如下：

• 具体来说，首先取数组中前k个字符，保存到堆中，顺序堆会自动调整。
• 然后从k+1开始遍历数组，每次都和堆顶元素进行比较。如果我们要求第k大的数，那么需要建立小顶堆。那么遍历的元素如果比堆顶元素小，那么堆顶弹出，把遍历元素放入堆中。
• 完全遍历之后，堆顶元素就是第k大的元素。

求第k个大的数，需要建立小顶堆。小顶堆的意思是堆顶元素最小，那么在这个有k个元素的堆中，k-1个元素都比堆顶元素大，那么堆顶元素不就是第k大的了么。

同理如果求第k小的数，需要建立大顶堆。在堆中，k-1个元素都比堆顶元素小，那么堆顶元素就是第k小的元素。

在Java中，我们使用优先队列来实现上述操作，因为优先队列的底层实现就是堆。

// 求第k小的数
private static int heap(int[] arr, int k) {
    // 创建一个大小为k的，大顶堆
    PriorityQueue<Integer> queue = new PriorityQueue<>(k, new Comparator<Integer>() {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    });
    // 循环整个数组，让前k个元素直接存入大顶堆中，从第k+1个元素开始需要和堆顶进行比较
    for (int i = 0; i < arr.length; i++) {
        if (queue.size() < k) {
            queue.offer(arr[i]);
        } else {
            int top = queue.peek();
            if (top > arr[i]) {
                queue.poll();
                queue.offer(arr[i]);
            }
        }
    }
    return queue.poll();
}

这里我们需要注意的是，优先队列默认的是小顶堆，如果我们要创建大顶堆，需要重写比较器。

4. 对比直接排序方法和使用堆排序方法的耗时

下面进行了对比实验，对比了直接排序法和使用堆排序的耗时

public static void main(String[] args) {
      int k = 25;
      int[] arr = new int[200000000];
      Random random = new Random();
      for (int i = arr.length - 1, j = 0; i >= 0; i--) {
          arr[i] = random.nextInt();
      }
      Long start = System.currentTimeMillis();
      int sort = sort(arr, k); // 直接排序
      Long end = System.currentTimeMillis();
      System.out.print("使用普通排序耗时:");
      System.out.println(end - start);
      System.out.println(sort);

      Long start1 = System.currentTimeMillis();
      int heap = heap(arr, k);
      Long end1 = System.currentTimeMillis();
      System.out.print("使用heap排序耗时:");
      System.out.println(end1 - start1);
      System.out.println(heap);
  }

private static int sort(int[] arr, int k) {
    Arrays.sort(arr);
    return arr[k - 1];
}

使用普通排序耗时:24185ms
第k小的数为：-2147482892
使用heap排序耗时:528ms
第k小的数为：-2147482892

从结果可以看出来差距还是和明显的

这里需要注意判断是否需要出堆顶元素的条件

• 对于大顶堆来说 top > arr[i]
• 对于小顶堆来说 top < arr[i]

5. 图解举例说明具体过程

下面举一个例子，求第三小的数，图解如下：