堆排序应用之topK问题

Posted 2021-05-01 互联网之瞳

题目：求海量数据（正整数）按逆序排列的前k个数（topK），因为数据量太大，不能全部存储在内存中，只能一个一个地从磁盘或者网络上读取数据，请设计一个高效的算法来解决这个问题。 第一行用户输入K，代表要求得topK  随后的N（不限制）行，每一行是一个整数代表用户输入的数据 直到用户输入-1代表输入终止 请输出topK，空格分割。

思路：先开辟一个K大小的数组arr，然后读取K个数据存储到数组arr，读到K+1的时候，如果arr[K+1]小于arr中最小的值，那么就丢掉不管，如果arr[K+1]大于arr中最小的值，那么就把arr[K+1]和数组中最小的值进行交换，然后再读取K+2个数。这样就能解决这个问题。但是这个算法复杂度为K+(N-K)*K，K可以忽略不计，所以时间复杂度为O(KN)。那这个代码很容易就写出来。假如题目要求用到NlgK的时间复杂度，那么这里就需要使用堆这种数据结构来解决问题，而且还是小顶堆。具体思想还是和数组一样。

代码：

import java.util.Arrays;
import java.util.Scanner;

public class TopK {

    static int[] heap;
    static int index = 0;
    static int k;
    
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        k = scanner.nextInt();
        heap = new int[k];
        int x = scanner.nextInt();
        while(x!=-1){
            deal(x);  // 处理x
            x = scanner.nextInt();
        }
        System.out.println(Arrays.toString(heap));
    }

    /**
 * 如果数据量小于等于k，直接加入堆中
 * 等于k的时候，进行堆化
 * @param x
 */
    private static void deal(int x) {
        if (index<k) {
            heap[index++] = x;
            if (index==k) {
                // 堆化
                makeMinHeap(heap);
                System.out.println(Arrays.toString(heap));
            }
        }else if (heap[0]<x) {
            heap[0] = x;
            MinHeapFixDown(heap, 0, k);
            System.out.println(Arrays.toString(heap));
        }else {
            System.out.println(Arrays.toString(heap));
        }
    }
    static void makeMinHeap(int[] A){
        int n = A.length;
        for(int i = n/2-1;i>=0;i--){
            MinHeapFixDown(A,i,n);
        }
    }
    
    private static void MinHeapFixDown(int[] A, int i, int n) {
        // 找到左右孩子
        int left = 2 * i + 1;
        int right = 2 * i + 2 ;
        // 左孩子已经越界，i就是叶子节点
        if (left>=n) {
            return ;
        }
        // min 指向了左右孩子中较小的那个
        int min = left;
        if (right>=n) {
            min = left;
        }else {
            if (A[right]<A[left]) {
                min = right;
            }
        }
        // 如果A[i]比两个孩子都要小，不用调整
        if (A[i]<=A[min]) {
            return ;
        }
        // 否则，找到两个孩子中较小的，和i交换
        int temp = A[i];
        A[i] = A[min];
        A[min] = temp;
        // 小孩子那个位置的值发生了变化，i变更为小孩子那个位置，递归调整
        MinHeapFixDown(A, min, n);
    }

}

结果：

　　

总结：partition和堆都能解决顺序统计量问题，堆更适用于海量数据流，适用于不能全部存储在内存中的数据，相当于实时数据流，而partition适用于能够一次加载到内存的数组。