TopK问题，剑指offer40题

最小的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

方案

TopK问题主要有一下几种方法

1. sort排序后直接输出前k个
2. 如果数组元素值不是很大，可以考虑使用计数排序
3. 堆排序
4. 快速排序

方案一

sort后直接输出前k个
这种方法太无脑，不考虑

方案二

计数排序，只适用于数组元素最大值不是特别大的情况

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> res;
        // 计数排序
        vector<int> vi(10010, 0);
        for(int i = 0; i < arr.size(); i++) {
            vi[arr[i]]++;
        }
        for(int i = 0; i < vi.size(); i++) {
            if(res.size() == k) break;
            if(vi[i] > 0) {
                while(vi[i]-- && res.size() < k) {
                    res.push_back(i);
                }
            }
        }
        return res;
    }
};

方案三

堆排序
用一个堆来存放数组前k个元素，然后从第k+1个元素开始遍历，去替换掉堆顶元素（即堆里的最大值），这样就能保证数组遍历完之后堆里的k个元素一定是最小的，且能保证有序

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> res;
        // 堆排 大堆存前k个数，然后从k+1个元素开始遍历，如果小于堆顶元素，就把堆顶元素换出，
        // 就是不断地把这k个数中的大数换出来
        if(!k) return res;
        priority_queue<int> q;
        for(int i = 0; i < k; i++) q.push(arr[i]);
        for(int i = k; i < arr.size(); i++) {
            if(q.top() > arr[i]) {
                q.pop();
                q.push(arr[i]);
            }
        }
        for(int i = 0; i < k; i++) {
            res.push_back(q.top());
            q.pop();
        }
        return res;
    }
};

方案四

快排的思想，但是每次只排序左边的，不排序右边的

我们的目的是寻找最小的 k 个数。假设经过一次 partition 操作，枢纽元素位于下标 m，也就是说，左侧的数组有 m 个元素，是原数组中最小的 m 个数。那么：

• 若 k = m，我们就找到了最小的 k 个数，就是左侧的数组；
• 若 k<m ，则最小的 k 个数一定都在左侧数组中，我们只需要对左侧数组递归地 parition 即可；
• 若 k>m，则左侧数组中的 m 个数都属于最小的 k 个数，我们还需要在右侧数组中寻找最小的 k-m 个数，对右侧数组递归地 partition 即可。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> res;
        // 快排
        if(!k) return res;
        else if(arr.size() <= k) return arr;
        
        quick_select(arr, 0, arr.size()-1, k-1);
        
        for(int i = 0; i < k; i++) res.push_back(arr[i]);
        
        return res;
    }
    void quick_select(vector<int> &arr, int l, int r, int k) {
        int m = partition(arr, l, r);
        if(k == m) return ;
        else if(k < m) {
            quick_select(arr, l, m-1, k);
        } else {
            quick_select(arr, m+1, r, k);
        }
    }
    
    int partition(vector<int> &arr, int l, int r) {
        int temp = arr[l];
        while(l < r) {
            while(l < r && arr[r] > temp) r--;
            arr[l] = arr[r];
            while(l < r && arr[l] <= temp) l++;
            arr[r] = arr[l];
        }
        arr[l] = temp;
        return l;
    }
};