AcWing 53. 最小的k个数（堆+快速排序的轴点选择函数）

1.基于堆的实现

创建一个最大堆，控制其最大规模为k。若堆的规模小于k则不断将数组元素压入到堆中，否则则要根据当前元素和堆顶元素的大小来判断是否进行“交换”操作（堆顶元素弹出，当前数组元素压入）。不断重复此操作就可以将原数组中所有最小的k个数全部保留到堆中了。

具体代码如下：

class Solution {
public:
    vector<int> getLeastNumbers_Solution(vector<int> input, int k) {
        priority_queue<int, vector<int>, less<int>> q;
        vector<int> ans;
        for(auto e: input){
            if(q.size() < k) q.push(e);
            else if(q.top() > e){
                q.pop();
                q.push(e);
            }
        }
        while(!q.empty()){
            ans.push_back(q.top());
            q.pop();
        }
        return vector<int> (ans.rbegin(), ans.rend());
    }
};

该算法的时间复杂度是O(nlogk)，虽然略高，但是此算法既不会改变原数组，同时又很适合用于大量数据的情况。

2.基于快速排序的轴点选择函数的实现

参考我的上一篇博文AcWing 52. 数组中出现次数超过一半的数字（基于中位数和众数的线性复杂度的解法）中的第1部分。这个方法仍然要用到快速排序的轴点选择函数。我们令此函数每次都返回轴点的位置，如果这个位置恰好为k，则说明从数组开始到轴点位置的元素全都是我们要找的最小的k个数。反之则要根据轴点的位置调整数组区间大小进行进一步地划分，直至返回的轴点位置为k为止。在提取数组前k个元素后要注意进行排序。

具体代码如下：

class Solution {
public:
    int partition(vector<int> &nums, int lo, int hi){
        int i = lo, j = hi;
        while(i < j){
            while(nums[i] < nums[lo]) i++;
            while(nums[j] >= nums[lo]) j--;
            if(i < j) swap(nums[i], nums[j]);
        }
        swap(nums[i], nums[lo]);
        return i;
    }

    vector<int> getLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans;
        if(input.empty()) return ans;
        int lo = 0, hi = input.size() - 1;
        int index = partition(input, lo, hi);
        while(index != k - 1){
            if(index > k - 1){
                hi = index - 1;
                index = partition(input, lo, hi);
            }
            else{
                lo = index + 1;
                index = partition(input, lo, hi);
            }
        }
        for(int i = 0; i < k; i++){
            ans.push_back(input[i]);
        }
        sort(ans.begin(), ans.end());
        return ans;
    }
};

该算法的时间复杂度为O(n)，相对于上一个方法要低一些，不过此算法需要改变原数组同时又不太适用于大量数据的情况。