leetcode 剑指offer面试题40.最小的k个数

leetcode lcof面试题40.最小的k个数

原题链接

题目

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

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

0 <= k <= arr.length <= 10000
0 <= arr[i] <= 10000

解题思路：最大堆

这题可以用最大堆的思想，维护一个长度为k的最大堆数组，若新元素比最大堆堆顶元素大，则将当前新元素与堆顶元素交换，然后进行堆的维护（即最大堆的性质维护），维护过程为：从堆顶开始，若子节点存在比当前元素大的值，将其子节点与当前元素交换，再将当前元素索引赋值为交换子节点索引，直到遇到叶子节点或是没有子节点比当前元素大则循环结束。

vector<int> getLeastNumbers(vector<int>& arr, int k) {
    //判断输入数组为空or k==0情况
    if(k==0||arr.empty()) return vector<int>();
    //最大堆
    vector<int> ans(k,INT_MAX);
    for(int x:arr){
        //当前元素比堆顶小才有能进堆，否则比最大堆堆顶还大肯定不是前k小
        if(x<ans[0]){
            //堆顶更新
            ans[0]=x;
            //维护最大堆性质
            maintain_heap(ans);
        }

    }
    return ans;
}

//维护最大堆
void maintain_heap(vector<int>& ans){
    int i=0;
    while(i<ans.size()){
        //该节点没有左子节点,也就肯定没有右子节点，说明到叶子，直接退出。
        if((2*i+1)>=ans.size()) break;
        
        int left=ans[2*i+1];
        //max_num为子节点中最大值
        int max_num=left;
        //如果有右子节点
        if((2*i+2)<ans.size()){
            int right=ans[2*i+2];
            max_num=left>right?left:right;
            //要交换
            if(ans[i]<max_num){
                if(left>right){
                    //交换左节点和当前根
                    swap(ans[i],ans[2*i+1]);
                    i=2*i+1;
                }else{
                    swap(ans[i],ans[2*i+2]);
                    i=2*i+2;
                }
            }
            else{//不要交换，意味着已经平衡
                break;
            }
        }
        else{//只有左子节点
            if(ans[i]<max_num){//比左子节点小，交换
                swap(ans[i],ans[2*i+1]);
                i=2*i+1;
            }
            else break;//没比子节点小，结束
        }
        
    }
}

void swap(int& x,int& y){
    int t=x;
    x=y;
    y=t;
}

除了用自己手写的最大堆以外，cpp的stl容器priority_queue优先队列也可以快速完成最大堆需求。

算法时间复杂度O(nlogk)，空间复杂度O(k)。