[优先队列][quick select] LeetCode 面试题40. 最小的k个数

[优先队列][quick select] LeetCode 面试题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个后再排个序，然后就不会了，

智障啊，不就是优先队列嘛，好几个月不碰啥都忘了

先把前k个数插入进去，然后再依此对比，如果准备入堆的数小于堆顶就弹出堆顶O(nlogk)

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(k == 0) return vector<int> {};
        priority_queue<int> q;//默认大根堆
        for(int i = 0; i < k; ++i) q.push(arr[i]);
        int n = arr.size();
        for(int i = k; i < n; ++i){
            if(arr[i] < q.top()) {
                q.pop();
                q.push(arr[i]);
            }
        }
        vector<int> res(k);
        for(int i = 0; i < k; ++i) {res[i] = q.top(); q.pop();}
        return res;
    }
};

看到官方还有python的堆，那是真滴流批

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0: return []
        hp = [-x for x in arr[:k]]	#插入前k个数，默认是小根堆，所以插入相反数
        heapq.heapify(hp)			#用列表初始化堆
        for idx,num in enumerate(arr[k:]):
            if num < -hp[0]:
                heapq.heappop(hp)
                heapq.heappush(hp,-num)
        res = [-x for x in hp]
        return res

上面那不是流批的就是摸拟，下面这个是真滴流批

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        return heapq.nsmallest(k,arr)	#返回前k个最小的数

quick select/快速选择

利用快速排序的特性

快排首先找一个基础元素t，然后将数组划分成左边小于t，右边大于t

如果t的下标等于k，那么t左边的元素就是前k小的元素

如果小于k，就继续划分右边的区间

大于k就划分左边的区间

和快排的区别就在于快选每次只用处理一边的元素

注意在选择时选择的是数组的下标，所以传输参数时要将k减去一，在返回数组时把1加回来

期望时间复杂度为 O(n)，最坏情况下的时间复杂度为 O(n^2)。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(k == 0) return vector<int> {};
        return quickSearch(arr,0,arr.size() - 1, k - 1);//注意减去1
    }

    vector<int> quickSearch(vector<int>& arr, int lo, int hi, int k){
        int idx = partition(arr,lo,hi);
        if(idx == k) return vector<int> (arr.begin(), arr.begin() + k + 1);//注意将1加回
        return idx > k ? quickSearch(arr, lo, idx - 1, k) : quickSearch(arr, idx + 1, hi, k);
    }

    int partition(vector<int>& arr, int lo, int hi){
        int t = arr[lo], idx = lo;
        while(lo<hi){
            while(lo < hi && arr[hi] >= t) --hi;
            while(lo < hi && arr[lo] <= t) ++lo;
            if(lo>=hi) break;
            int temp = arr[hi];
            arr[hi] = arr[lo];
            arr[lo] = temp;
        }
        arr[idx] = arr[lo];
        arr[lo]  = t;
        return lo;
    }
};

计数排序

因为本题的数据范围有限，所以用计数排序也可以，时间复杂度O(n)

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(k == 0) return vector<int> {};
        vector<int> count(10005,0);
        for(auto num:arr) ++count[num];
        vector<int> res(k);
        int t = 0;
        for(int i = 0; i < 10005; ++i) {
            while(count[i]-- > 0){
                if(t == k) break;
                res[t++] = i;
            }
        }
        return res;
    }
};