堆排序在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;    }};