【牛客题霸题解】最小的K个数

如果对整个数组进行排序的话，前 $k$ 个元素就是最小的 $k$ 个数，但这样做时间成本就比较高了。
我们来介绍一下这道题目的两种比较优秀的解法。一种是利用了快排的部分思路，另一种利用了堆排的部分思路。
先介绍一下快排相关的解法。利用快排的partition函数来进行划分。如果我们的基准数划分完成了以后在第 $k$ 位上的话，那基准数以及基准数前面的元素就是最小的 $k$ 个数了。

class Solution {
public:
    //快排思路
    int partition(vector<int>&arr,int l,int r)
    {
        int x = arr[l];
        while( l < r )
        {
            while(l < r && arr[r] >= x) { r--;}
            arr[l] = arr[r];
            while(l < r && arr[l] < x) { l++;}
            arr[r] = arr[l];
        }
        arr[l] = x;
        return l;
    }
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans;
        int n = input.size();
        if(k > n ||k<=0){return ans;}
        int l = 0;
        int r = n-1;
        int mid = partition(input,l,r);
        while(mid != k-1)
        {
            if(mid > k-1)
            {
                r = mid-1;
                mid = partition(input,l,r);
            }
            else{
                l = mid+1;
                mid = partition(input,l,r);
            }
        }
        for(int i = 0 ; i < k ; i++)
        {
            ans.push_back(input[i]);
        }
        return ans;

    }
};

这个题目可以借助数据结构堆在 $O(nlogk)$ 的时间复杂度完成。并且不会修改原数组，特别适合处理海量的数据且k比较小的时候来进行。下面我们来介绍一下做法。
创建一个大根堆，限定其最多只能放 $k$ 个元素。然后依次取待排序的数据 $arr\left[i\right]$ 往堆里放，堆没满的话直接插入就可以了。如果堆满了的话，那目前堆顶元素就是堆中的最大值。这个时候我们来比较一下 $arr\left[i\right]$ 和堆顶元素的大小。如果 $arr\left[i\right]$ 比堆顶元素大的话，那就说明数据中已经有 $k$ 个元素比 $arr\left[i\right]$ 小了。 $arr\left[i\right]$ 已经不可能属于最小的 $k$ 个数了。那就直接忽略掉 $arr\left[i\right]$ 继续看下一个数。如果 $arr\left[i\right]$ 比堆顶元素小的话，那就将堆顶元素删除，把 $arr\left[i\right]$ 插入到堆中。
堆在C++中的优先队列底层就是用堆实现的。Java和python中也可以使用优先队列。这样的话，我们直接利用优先队列就可以实现 $O\left(1\right)$ 的时间复杂度获取堆顶元素， $O\left(logn\right)$ 的复杂度进行插入和删除了

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans ;
        int N = input.size();
        if(k > N||k==0){return ans;}
        priority_queue<int> heap;
        for(int i = 0 ; i < N ; i++)
        {
            if(i < k){heap.push(input[i]);}
            else{
                if(heap.top() > input[i])
                {
                    heap.pop();
                    heap.push(input[i]);
                }
            }
        }
        while(!heap.empty())
        {
            ans.push_back(heap.top());
            heap.pop();
        }
        return ans;
    }
};

下面再附上一份手写堆的代码。

class Solution {
public:
    //堆排思路（手写堆
    int heapSize;//记录堆的大小
    int Left(int i){
         if(i*2+1>=heapSize) return -1;
         return i*2+1;
    }
    int Right(int i){
         if(i*2+2>=heapSize) return -1;
         return i*2+2;
    }
    void minHeapify(vector<int> & heap,int i) {
        int L = Left(i);
        int R = Right(i);
        int maxi = i;
        if(L != -1 && heap[L]>heap[i]){ maxi = L;}
        if(R != -1 && heap[R]>heap[maxi]){ maxi = R;}
        if(maxi != i){
            swap(heap[i],heap[maxi]);
            i = maxi;
            minHeapify(heap,i);
        }
    }
    void Delete(vector<int> & heap) {
        swap(heap[heapSize-1],heap[0]);
        heapSize--;
        minHeapify(heap,0);
    }

    void Insert(vector<int> & heap,int n)
    {
        heap[heapSize] = n;
        int i = heapSize;
        heapSize++;
        while(i!=0&& heap[i]>heap[(i-1)/2])
        {
            swap(heap[i],heap[(i-1)/2]);
            i = (i-1)/2;
        }
    }
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans;
        int n = input.size();
        if(k>n || k<=0 ){return ans;}
        heapSize = 0;
        vector<int> heap(k);
        for(int i = 0 ; i < n ; i++) {
            if(heapSize < k) {Insert(heap,input[i]);}
            else{
                if(heap[0] > input[i]){
                    Delete(heap);
                    Insert(heap,input[i]);
                }
            }
        }
        return heap;
    }
};