堆與堆排序與topK問題

電面的時候問了經典的topK問題，沒準備到被問了個措不及防，現在把相關知識點記錄下來。
假設我們有一些數據，需要按照某種關鍵字求出這些數據中最小(大)的K個數，即求出這些數據的topK。
當數據量較小的時候，一個簡單的想法是直接對數據進行排序，然後取最前面的K個數據；但是當數據量較大，數據無法一次性放入內存的時候應該怎麼辦呢？
這時候就需要藉助堆這種數據結構。堆通常是一個可以被看做一棵樹的數組對象，其總是滿足下列性質：
(a)堆中某個節點的值總是不大於或不小於其父節點的值
(b)堆總是一棵完全二叉樹
節點值不大於其父節點的堆稱爲大根堆，反之稱爲小根堆。

堆排序與topK問題無關，以下內容純屬展開：
以堆爲基礎的排序方法稱爲堆排序。堆排序主要可以分爲以下幾個步驟：
(a)建堆
(b)將堆頂與堆的最後一個元素對換
(c)將堆的大小-1(此時最後一個元素看作已排序好)，並對剩下的部分保持堆的性質
(d)重複(b)、(c)直至排序完成
代碼如下：

public class MinHeap {
    public final static int MAX = 100;
    private int[] heap = null;
    private int size = 0;
    public MinHeap(){
        heap = new int[MAX];
    }
    public boolean add(int num){
        if (size >= MAX){
            return false;
        }
        heap[size] = num;
        size++;
        return true;
    }
    public void buildHeap(){
        for (int i = (size - 1) / 2;i >= 0;i--){
            heapify(i);
        }
    }
    private void heapify(int pos){
        int left = pos * 2 + 1, right = pos * 2 + 2;
        if (left >= size){
            return;
        }else if (right >= size){
            if (heap[pos] > heap[left]){
                int tmp = heap[pos];
                heap[pos] = heap[left];
                heap[left] = tmp;
            }
            return;
        }else{
            int minPos = pos;
            if (heap[pos] > heap[left]){
                if (heap[left] > heap[right]){
                    minPos = right;
                }else{
                    minPos = left;
                }
            }else if (heap[pos] > heap[right]){
                minPos = right;
            }
            int tmp = heap[pos];
            heap[pos] = heap[minPos];
            heap[minPos] = tmp;
            if (minPos != pos){
                heapify(minPos);
            }
        }
    }
    public void heapSort(){
        buildHeap();
        int length = size;
        while (size > 0){
            int tmp = heap[size-1];
            heap[size-1] = heap[0];
            heap[0] = tmp;
            size--;
            heapify(0);
        }
        size = length;
    }
    public void print(){
        for (int i = 0;i < size;i++){
            System.out.print(heap[i] + " ");
        }
        System.out.println();
    }
    public static void main(String[] args) {
        MinHeap heap = new MinHeap();
        for (int i = 10;i > 0;i--){
            heap.add(i);
        }
        heap.heapSort();
        heap.print();
    }
}

堆排序heapSort()的時間複雜度爲O(nlgn)，其中建堆buildHeap()的時間複雜度爲O(n)(不是筆誤，具體見算法導論)，每一輪保持堆的性質heapify(int pos)的時間複雜度爲O(lgn)，一共O(n)輪。

回到topK問題上，我們可以維護一個大小爲K的堆來幫助我們解決topK問題。以最小的K個數據爲例，我們維護一個大小爲K的大根堆在內存中，接下來每次從那一大堆數據當中讀出一個數據與堆頂比較，若該數據比堆頂數據小，則說明它比當前topK的最大值要小，因此將堆頂替換爲該數據，再用heapify(0)保持該堆的最大堆性質即可。