Leetcode 剑指 Offer 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]

方法一：最直接的思路，简单题就要对应简单做法。用sort函数对 arr 排序，再取其前 k 个数即可。先放C++代码，思路简单易懂。

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
void print(vector<int> inputvector){
    for(int i:inputvector){
        cout<<i<<endl;
    }
}
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(arr.empty() || k == 0) {return {};}
        vector<int> res;
        sort(arr.begin(), arr.end());
        for(int i : arr) {
            if(k <= 0) {break;}
            res.push_back(i);
            -- k;
        }
        return res;
    }
};

int main()
{   
    vector<int> myvector{4,5,1,6,2,7,3,8};
    Solution S;
    print(S.getLeastNumbers(myvector,4));
    system("pause");
    return 0;
}

方法二：利用最大堆（priority_queue）。

1. 先将 arr 的前 k 个数入堆，因为是最大堆，所以堆顶的数一定是最大的。
2. 那么我们再从 arr 的第 k + 1 个元素迭代起，凡是遇到比当前堆顶元素小的，就将该元素 push 入堆并 pop 掉堆顶元素。
3. 注意，每次执行完第二步之后，堆会自动更新，以保证堆顶元素始终为堆中所有元素的最大值。
   先放C++代码，思路清晰明了。

#include <iostream>
#include <vector>
#include<queue>
using namespace std;
void print(vector<int> inputvector){
    for(int i:inputvector){
        cout<<i<<endl;
    }
}
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(arr.empty() || k == 0) {return {};}
        vector<int> res(k);
        priority_queue<int> max_heap;

        for(int i = 0; i < k; ++i) {max_heap.push(arr[i]);} // 用 arr 的前 k 个数填充最大堆
        for(int i = k; i < arr.size(); ++i) {
            if(arr[i] < max_heap.top()){
                max_heap.pop();
                max_heap.push(arr[i]); // 循环更新最大堆
            }
        }
        for(int i = 0; i < k; ++i) {
            res[i] = max_heap.top(); // 填充 res
            max_heap.pop();
        }

        return res;
    }
};

int main()
{   
    vector<int> myvector{4,5,1,6,2,7,3,8};
    Solution S;
    print(S.getLeastNumbers(myvector,4));
    system("pause");
    return 0;
}

方法三：快速选择——退化的快速排序。先放C++代码，思路清晰明了，注释都已写好在代码中。

#include <iostream>
#include <vector>
#include<queue>
using namespace std;
void print(vector<int> inputvector){
    for(int i:inputvector){
        cout<<i<<endl;
    }
}
class Solution {
private:
    vector<int> res;
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(arr.empty() || k == 0) {return {};}
        return quickSelection(arr, 0, arr.size() - 1, k - 1); // 注意第 k 个数对应的下标是 k - 1
    }
    vector<int> quickSelection(vector<int>& arr, int left, int right, int index) {
        // partition函数将一个区间内所有小于下标为 j 的数放在 j 左边，大于下标为 j 的数放在 j 右边
        int j = partition(arr, left, right); 
        
        if(j == index) { // 若 j 刚好等于 k - 1，将 arr[0] 至 arr[j] 输入 res
            for(int i = 0; i < j + 1; ++i) {res.push_back(arr[i]);}
            return res;
        }
        // 若 j 小于 k - 1，将区间变成 [j + 1, right]；反之，区间变成 [left, j - 1]
        return j < index ? quickSelection(arr, j + 1, right, index) : quickSelection(arr, left, j - 1, index);
    }
    int partition(vector<int>& arr, int left, int right) {
        int value = arr[left];
        int i = left, j = right + 1;

        while(true) {
            while(++ i <= right && arr[i] < value); // 找到从左往右第一个大于等于 value 的下标
            while(-- j >= left && arr[j] > value); // 找到从右往左第一个小于等于 value 的下标
            if(i >= j) {break;} // 如果找不到，说明已经排好序了，break
            swap(arr[i], arr[j]); // 如果找到了，交换二者
        }
        swap(arr[left], arr[j]); // arr[j]是小于 value 的，这一步使得所有小于下标为 j 的数都在 j 左边

        return j;
    }
    void swap(int& a, int& b) { 
        int temp = a;
        a = b;
        b = temp;
    }
};

int main()
{   
    vector<int> myvector{4,5,1,6,2,7,3,8};
    Solution S;
    print(S.getLeastNumbers(myvector,4));
    system("pause");
    return 0;
}