堆排序在TOP K问题中的应用

客官°小女子只卖身不卖艺 2022-08-09 11:55 107阅读 0赞

# 问题 #

从数组中找出最大或者最小的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;
        }
    };