最小堆解決Top K問題

問題描述:

有一組數據n個,要求取出這組數據中最大的K個值。
對於這個問題,解法有不少中。好比排序及部分排序,不過效率最高的要數最小堆,它的時間複雜度爲O(nlogk)。java

解題思路:

  1. 取出數組的前n個元素,建立長度爲n的最小堆。
  2. 從n開始循環數組的剩餘元素,若是元素(a)比最小堆的根節點大,將a設置成最小堆的根節點,並讓堆保持最小堆的特性。
  3. 循環完成後,最小堆中的全部元素就是須要找的最大的n個元素。

Java實現:

1.建堆代碼:web

public class Heap<T> {  

    //以數組形式存儲堆元素
    private T[] heap;  

    //用於比較堆中的元素。c.compare(根,葉子) > 0
    //使用相反的Comparator能夠建立最大堆、最小堆
    private Comparator<? super T> comparator;  

    //初始化堆
    public Heap(T[] heap, Comparator<? super T> comparator) {  
        this.heap = heap;  
        this.comparator = comparator;  
        buildHeap();  
    }  

    //返回值爲i/2
    private int parent(int i) {  
        return (i - 1) >> 1;  
    }  

    //返回指定節點的left子節點數組索引。至關於2*(i+1)-1
    private int left(int i) {  
        return ((i + 1) << 1) - 1;  
    }  

    //返回指定節點的right子節點數組索引。至關於2*(i+1)
    private int right(int i) {  
        return (i + 1) << 1;  
    }  

    //堆化,i是堆化的起始節點
    private void heapify(int i) {  
        heapify(i, heap.length);  
    }  

    //堆化,i是堆化的起始節點,size是堆化的範圍
    private void heapify(int i, int size) {  
        int l = left(i);  
        int r = right(i);  
        int next = i;  
        if (l < size && comparator.compare(heap[l], heap[i]) > 0)  
            next = l;  
        if (r < size && comparator.compare(heap[r], heap[next]) > 0)  
            next = r;  
        if (i == next)  
            return;  
        swap(i, next);  
        heapify(next, size);  
    }  

    //對堆排序
    public void sort() {  
        // buildHeap(); 
        for (int i = heap.length - 1; i > 0; i--) {  
            swap(0, i);  
            heapify(0, i);  
        }  
    }   

    //交換數組值 
    private void swap(int i, int j) {  
        T tmp = heap[i];  
        heap[i] = heap[j];  
        heap[j] = tmp;  
    }  

    //建立堆
    private void buildHeap() {  
        for (int i = (heap.length) / 2 - 1; i >= 0; i--) {  
            heapify(i);  
        }  
    }  

    public void setRoot(T root) {  
        heap[0] = root;  
        heapify(0);  
    }  

    public T root() {  
        return heap[0];  
    }   
}

2.最小堆實現Top K篩選api

public class TopK {  

    public static int[] toPrimitive(Integer array[]) {  
        if (array == null)  
            return null;  
        if (array.length == 0)  
            return new int[0];  
        int result[] = new int[array.length];  
        for (int i = 0; i < array.length; i++){
            result[i] = array[i].intValue();
        }
        return result;  
    }  

    public static int[] topK(int[] array, int n) {  
        if (n >= array.length) {  
            return array;  
        }

        Integer[] topn = new Integer[n];  
        for (int i = 0; i < topn.length; i++) {  
            topn[i] = array[i];  
        }

        Heap<Integer> heap = new Heap<Integer>(topn, new Comparator<Integer>() {   
            @Override  
            public int compare(Integer o1, Integer o2) {  
                // 生成最大堆使用o1-o2,生成最小堆使用o2-o1 
                return o2 - o1;  
            }  
        });  
        for (int i = n; i < array.length; i++) {  
            int value = array[i];  
            int min = heap.root();  
            if (value > min) {  
                heap.setRoot(value);  
            }  
        }   
        return toPrimitive(topn);  
    }  

    public static void main(String[] args) {
        Random random = new Random(100);
        int[] array = new int[100];
        for(int i=0; i<100; i++){
            array[i] = random.nextInt(1000);
        }
        int count = 0;
        for(int i=0; i<100; i++){
            count++;
            System.out.print(array[i] + " ");
            if(count >= 25){
                System.out.println();
                count=0;
            }
        }
        System.out.print("Top 10 :");
        int[] result = topK(array, 10);
        for (int integer : result) {  
            System.out.print(integer + ",");  
        }  
    }  

}