算法-Top K问题

问题描述

现有一组千万级别的数，找出其中最大（小）的K个数。

思路

利用大（小）顶堆所有子节点元素都比父节点元素小（大）的性质来实现的。
故可先建立一个包含K的元素的堆，然后便利集合，如果集合元素比堆顶大（小），那就用该元素来替换堆顶元素，同时维护堆的性质，最后遍历结束之后，堆中元素即为最大（小）的K的数。

代码

public class TopK {
    public static void main(String[] args){
        int[] a = { 1, 17, 3, 4, 5, 6, 7, 16, 9, 10, 11, 12, 13, 14, 15, 8 };
        int[] b = topK(a,4);
        for(int tmp : b){
            System.out.print(tmp+",");
        }
    }
    
    public static int[] topK(int[] arr,int k){
        int[] head = new int[k];
        for(int i = 0;i<k;i++){
            head[i] = arr[i];
        }
        buildHeap(head,k);
        for(int i = k;i<arr.length;i++){
            if(arr[i] < head[0]){
                setTop(head,arr[i],k);
            }
        }
        return head;
    }
    
    public static void setTop(int[] arr,int top,int length){
        arr[0] = top;
        sink(arr,0,length);
    }
    
    public static void sink(int[] arr,int k,int length){
        while(k <= length / 2){
            int j = k * 2;
            if(j < length && arr[j] < arr[j+1]){
                j++;
            }
            if(arr[k] < arr[j]){
                swap(arr,k,j);
                k = j;
            }else{
                break;
            }
        }
    }
    
    public static void buildHeap(int[] arr,int length){
        for(int i = (length-1) / 2;i >= 0; i--){
            sink(arr,i,length);
        }
    }
    
    
    public static void swap(int[] arr,int a,int b){
        int tmp = arr[a];
        arr[a] = arr[b];
        arr[b] = tmp;
    }
}