1. 程式人生 > >最小堆解決Top K問題

最小堆解決Top K問題

問題描述:

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

解題思路:

  1. 取出陣列的前n個元素,建立長度為n的最小堆。
  2. 從n開始迴圈陣列的剩餘元素,如果元素(a)比最小堆的根節點大,將a設定成最小堆的根節點,並讓堆保持最小堆的特性。
  3. 迴圈完成後,最小堆中的所有元素就是需要找的最大的n個元素。

Java實現:

1.建堆程式碼:

 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篩選

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 + ",");  
        }  
    }  

}