TOP-K 堆实现

题目描述

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

题解

TOP-K问题使用堆实现可以实现最小的时间复杂度。方法是先取出数组中的k个数构造大顶堆（最小的K个值），然后逐个读取数组中的剩下的数，并与堆顶比较，如果小于堆顶，则与堆顶元素替换，再调整堆。数组读取完成后，输出堆中的k个元素就是最小的k个数字。

使用STL自己实现堆结构

class Solution {
public:
    void adjust(vector<int> &v, int start, int len){
        int tmp = v[start];
        for(int i = 2*start+1; i < len; i = i*2+1){ // i*2+1为左子节点，i*2+2为右子节点
            if(i != (len-1) && v[i] < v[i+1])  // 如果右节点大于左节点
                i++;
            if(tmp < v[i]){
                v[start] = v[i];
                start = i;
            }else
               break;
        }
        v[start] = tmp;
    }
    void buildHeap(vector<int> &v){
        for(int i = v.size()/2-1; i>=0; i--)  // 从第一个非叶结点开始建堆
            adjust(v, i, v.size());
    }
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        if(input.size()==0)
            return input;
        if(k>input.size() || k <=0){
            vector<int> v;
            return v;
        }
        buildHeap(input);
        for(int i = k; i < input.size(); i++){
            if(input[i] < input[0]){
                int tmp = input[0];
                input[0] = input[i];
                input[i] = tmp;
                adjust(input, 0, k);  //调整大小为k的堆
            }
        }
        vector<int> v(input.begin(), input.begin()+k);
        return v;
    }
};

STL的priority_queue已经使用了大顶堆进行实现，所以可以直接使用priority_queue实现。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        if(input.size()==0)
            return input;
        if(k>input.size() || k <=0){
            vector<int> v;
            return v;
        }
        priority_queue<int> q;
        for(int i = 0; i < input.size(); i++){
            if(q.size()<k)
                q.push(input[i]);
            else if(input[i] < q.top()){
                q.pop();
                q.push(input[i]);
            }
        }
        vector<int> v;
        while(!q.empty()){
            v.push_back(q.top());
            q.pop();
        }
        return v;
    }
};

若想使用priority_queue来实现查找最大的K个数，则需要使用小顶堆，stl中声明小顶堆的priority_queue的方法如下：

priority_queue<int, vector<int>, greater<int>> q;