如何找出排名前 500 的数？

题目描述

有 20 个数组，每个数组有 500 个元素，并且有序排列。如何在这 20*500 个数中找出前 500 的数？

解答思路

对于 TopK 问题，最常用的方法是使用堆排序。对本题而言，假设数组降序排列，可以采用以下方法：

首先建立大顶堆，堆的大小为数组的个数，即为 20，把每个数组最大的值存到堆中。

接着删除堆顶元素，保存到另一个大小为 500 的数组中，然后向大顶堆插入删除的元素所在数组的下一个元素。

重复上面的步骤，直到删除完第 500 个元素，也即找出了最大的前 500 个数。

为了在堆中取出一个数据后，能知道它是从哪个数组中取出的，从而可以从这个数组中取下一个值，可以把数组的指针存放到堆中，对这个指针提供比较大小的方法。

  1. import lombok.Data;
  3. import java.util.Arrays;
  4. import java.util.PriorityQueue;
  6. /**
  7.  * @author https://github.com/yanglbme
  8.  */
  9. @Data
  10. public class DataWithSource implements Comparable<DataWithSource> {
  11.     /**
  12.      * 数值
  13.      */
  14.     private int value;
  16.     /**
  17.      * 记录数值来源的数组
  18.      */
  19.     private int source;
  21.     /**
  22.      * 记录数值在数组中的索引
  23.      */
  24.     private int index;
  26.     public DataWithSource(int value, int source, int index) {
  27.         this.value = value;
  28.         this.source = source;
  29.         this.index = index;
  30.     }
  32.     /**
  33.      *
  34.      * 由于 PriorityQueue 使用小顶堆来实现，这里通过修改
  35.      * 两个整数的比较逻辑来让 PriorityQueue 变成大顶堆
  36.      */
  37.     @Override
  38.     public int compareTo(DataWithSource o) {
  39.         return Integer.compare(o.getValue(), this.value);
  40.     }
  41. }
  43. class Test {
  44.     public static int[] getTop(int[][] data) {
  45.         int rowSize = data.length;
  46.         int columnSize = data[0].length;
  48.         // 创建一个columnSize大小的数组，存放结果
  49.         int[] result = new int[columnSize];
  51.         PriorityQueue<DataWithSource> maxHeap = new PriorityQueue<>();
  52.         for (int i = 0; i < rowSize; ++i) {
  53.             // 将每个数组的最大一个元素放入堆中
  54.             DataWithSource d = new DataWithSource(data[i][0], i, 0);
  55.             maxHeap.add(d);
  56.         }
  58.         int num = 0;
  59.         while (num < columnSize) {
  60.             // 删除堆顶元素
  61.             DataWithSource d = maxHeap.poll();
  62.             result[num++] = d.getValue();
  63.             if (num >= columnSize) {
  64.                 break;
  65.             }
  67.             d.setValue(data[d.getSource()][d.getIndex() + 1]);
  68.             d.setIndex(d.getIndex() + 1);
  69.             maxHeap.add(d);
  70.         }
  71.         return result;
  73.     }
  75.     public static void main(String[] args) {
  76.         int[][] data = {
  77.                 {29, 17, 14, 2, 1},
  78.                 {19, 17, 16, 15, 6},
  79.                 {30, 25, 20, 14, 5},
  80.         };
  82.         int[] top = getTop(data);
  83.         System.out.println(Arrays.toString(top)); // [30, 29, 25, 20, 19]
  84.     }
  85. }

方法总结

求 TopK，不妨考虑一下堆排序？