海量数据-两种方法解决top k问题

假如提供一百万个数据(甚至更大)的无序数组,如何从中获取最大的k个元素?

最容易想到的是先降序排序然后获取前k个元素,那假设我们用最常用的O(nlogn)级别的排序算法,要获取top k,时间成本也就是排序成本了,为O(nlogn),而这里的n是百万级别的,那能不能更快一点呢?

最小堆

假设我们设定一个最小堆,并固定这个最小堆的最大大小为k,那么任何一时刻堆顶元素都是这个堆中的第k大的数据(也是最小的数据)。来一个数据就和堆顶元素比较,如果比堆顶元素大就替换堆顶元素,并堆化。最后当所有数据都走一遍后,堆内的元素就是所有数据中最大的k个元素,而堆顶元素就是第k大的元素。

那么我们来分析一下它在时间复杂度上优化了多少吧。首先所有数据都要走一遍比较,就是n,其次维护了一个大小为k的堆,也就是说堆化的时间复杂度为O(log k),那么总的时间复杂度就是O(nlogk) ,乍一看和先排序的时间复杂度O(nlogn)也没差多少,甚至说是一个衡量级别的,但毕竟这里说的是海量数据,n可能很大,k可能比较小,那么时间上就快很多了。

它的另一个好处是空间复杂度也仅为O(k),不需要额外维护一个O(n)级别的海量数据存储。

代码实现:

public static List<Integer> topKByHeap(int[] arr, int k) {
    if (arr == null || arr.length == 0 || k <= 0 || k > arr.length) {
        return new ArrayList<>();
    }
    PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    for (int i = 0; i < k; i++) {
        minHeap.add(arr[i]);
    }
    for (int i = k; i < arr.length; i++) {
        if (arr[i] > minHeap.peek()) {
            minHeap.poll();
            minHeap.add(arr[i]);
        }
    }
    return new ArrayList<>(minHeap);
}

快排思想

海量数据-两种方法解决top k问题

O(nlogk)好像还是有点不够看,能不能直接优化成O(n)呢?

我们知道快排的思想是每次随机选取一个pivot把数组划分成相对有序的两部分,那假设第一次随机选取pivot并分区之后,后面这部分的元素数量就正好为k,不就说明后面这部分数据就是最大的k个元素,那最好时间复杂度岂不是O(logn)

海量数据-两种方法解决top k问题

那么问题就成为了选中一个pivot,将数组分为前后两部分,后面部分元素的数量为k的问题。

虽然很难一次就中,但我们在一次分区后,通过比较k和后面部分元素的数量就知道了,后面部分的元素是分多了还是分少了。

如果后面的元素数量分多了,就继续在后面这部分数据中随机选取pivot,重复步骤。如果后面的元素分少了,则往前选取pivot,就形成了一种类似二分查找的步骤。

那么我们来分析一下它的时间复杂度,我们知道快速排序每次分区的时间复杂度就是log(n),那么我们假设每次分区前后两部分的数量相近,那么时间复杂度就是log(n)+log(n/2)+log(n/4)+...,取极限为log(2n),即log(n)级别,是不是比一开始的log(nlogn)快很多?特别是在海量数据的情况下。

代码实现:

public static List<Integer> topK(int[] arr, int k) {
    if (arr == null || arr.length == 0 || k <= 0 || k > arr.length) {
        return new ArrayList<>();
    }
    // 如果是 BottomK 的话就不需要单独处理 k==n 情况了,因为划分是小于等于放在左
    // topK 的话如果是将大于的 k 个放在右边,那么 pivot 就变成了-1,会形成死循环。
    if (k == arr.length) {
        return IntStream.of(arr).boxed().collect(Collectors.toList());
    }
    int low = 0;
    int high = arr.length - 1;
    int index = QuickSort.partition(arr, low, high);
    while (index != arr.length - k - 1) {
        if (index > arr.length - k - 1) {
            index = QuickSort.partition(arr, low, index - 1);
        } else {
            index = QuickSort.partition(arr, index + 1, high);
        }
    }
    List<Integer> list = new ArrayList<>(k);
    for (int i = index + 1; i < arr.length; i++) {
        list.add(arr[i]);
    }
    System.out.println(Arrays.toString(arr));
    return list;
}

这里的Quicksort得自己实现了,这里不作为重点,可以参考十大排序算法与Java实现

如果是找第K大的元素呢?

作为拓展,讨论一下如果要找第K大的数据,仍然可以用这两个方法:

  • 对于最小堆则是操作一遍之后直接取堆顶元素。
  • 对于快排思想,则是找一个pivot,使得分区之后它正好处于第arr.length-k+1的位置上。

原创文章,作者:彭晨涛,如若转载,请注明出处:https://www.codetool.top/article/%e6%b5%b7%e9%87%8f%e6%95%b0%e6%8d%ae-%e4%b8%a4%e7%a7%8d%e6%96%b9%e6%b3%95%e8%a7%a3%e5%86%b3top-k%e9%97%ae%e9%a2%98/