堆排序在TOP K問題中的應用

問題

從數組中找出最大或者最小的k個數。

思路

以最小的k個數爲例。可以使用一個大小爲k的數組，然後依次遍歷原始數據，當有元素比數組裏的元素小時，就用這個數據將其替換出來。思路是對的，但是從大小爲k的數組裏面搜索最大元素的複雜度是O(n)。接下來優化一下，我們知道堆排序獲得最大值（最小值）的複雜度是O(1)，調整堆的複雜度是O(log n)。在海量數據處理的時候這個優化的效果是很明顯的。

代碼

題目描述：
輸入n個整數，找出其中最小的K個數。

class BigHeap{//大頂堆，用於從k個數中獲取最大值。
private:
    vector<int> data;//堆是完全二叉樹，用數組存放
    int len;//堆中元素的個數
public:
    BigHeap(int n):len(n){}
    void push(int elem){//把元素放入堆中，然後調整堆
        int size = data.size();
        if(len < size)
            data.at(len) = elem;
        else
            data.push_back(elem);
        len += 1;
        int fa_idx = (len-2)/2;//新元素父節點下標
        int new_idx = len-1;//新元素的下標
        while(fa_idx >= 0){//每次講新元素放到尾部，然後向上調整，直到滿足條件或者到堆頂。
            if(data.at(new_idx) > data.at(fa_idx)){//因爲是大頂堆，所以比父親大就與父親交換
                int tmp = data.at(new_idx);
                data.at(new_idx) = data.at(fa_idx);
                data.at(fa_idx) = tmp;
                new_idx = fa_idx;
                fa_idx = (fa_idx-1)/2;
            }
            else{
                break;
            }
        }
        printHeap();
    }
    int getMax(){//堆頂元素就是最大值
        if(len > 0)
            return data.at(0);
        else
            return 100000;
    }
    void printHeap(){//調試使用
        for(int i = 0; i < len; i++)
            cout << data.at(i) << " ";
        cout << endl;
    }
    void deleteMax(){//刪除堆頂元素，然後調整堆
        if(len < 1)
            return;
        data.at(0) = data.at(len-1);//刪除堆頂元素的方法是與最後的元素交換，然後減小元素個數len
        len -= 1;
        int fa_idx = 0;
        int lch_idx = (0+1)*2 - 1;
        int rch_idx = (0+1)*2;
        while(lch_idx < len){//調整堆的方法是向下調整，直到滿足條件或者到堆尾
            int max_idx;
            if(rch_idx >= len){
                max_idx = lch_idx;
            }
            else{
                if(data.at(lch_idx) > data.at(rch_idx)){
                    max_idx = lch_idx;
                }
                else{
                    max_idx = rch_idx;
                }
            }
            if(data.at(fa_idx) < data.at(max_idx)){//向下調整是與孩子中最大的交換
                int tmp = data.at(fa_idx);
                data.at(fa_idx) = data.at(max_idx);
                data.at(max_idx) = tmp;
                fa_idx = max_idx;
                lch_idx = (max_idx + 1) * 2 -1;
                rch_idx = (max_idx + 1) * 2;
            }
            else{
                break;
            }
        }
        printHeap();
    }
};
class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        int len = input.size();
        if(k > len){
            vector<int> ans;
            return ans;
        }

        BigHeap myBigHeap(0);
        int i;
        for(i = 0; i < k; i++){//現將前k個元素放入堆中
            myBigHeap.push(input.at(i));
        }
        while(i < len){//如果新元素比堆中最大元素小，則刪除堆中最大元素，並將新元素入堆
            int max = myBigHeap.getMax();
            if(input.at(i) < max){
                myBigHeap.deleteMax();
                myBigHeap.push(input.at(i));
            }
            i++;
        }
        vector<int> ans(k, 0);//返回排序好的最小k個數
        for(i = k-1; i >= 0; i--){
            ans.at(i) = myBigHeap.getMax();
            myBigHeap.deleteMax();
        }
        return ans;
    }
};