剑指offer—topK问题，最小的k个数——手写大根堆（堆排变形），快排变形Partition

题意：输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

思路1

类似这种topK问题，都可以用堆排序的变形来解决
堆排序代码：

#include<iostream>
#include<vector>
using namespace std;

// 递归方式构建大根堆(len是arr的长度，index是第一个非叶子节点的下标)
void adjust(vector<int>& v, int len, int index){
    int left = 2*index + 1;
    int right = 2*index + 2;

    int maxIndex = index;
    if(left < len && v[left] > v[maxIndex]) maxIndex = left;
    if(right < len && v[right] > v[maxIndex]) maxIndex = right;

    if(maxIndex != index){
        swap(v[maxIndex], v[index]);
        adjust(v, len, maxIndex);
    }

}

void heapSort(vector<int>& v, int size)
{
    // 构建大根堆（从最后一个非叶子节点向上）
    //完全二叉树叶子节点数 = 总节点数 / 2，最后一个非叶子节点 = 总节点数/2 - 1
    for(int i = size/2 - 1; i >= 0; i--){
        adjust(v, size, i);
    }

    // 调整大根堆
    for(int i = size - 1; i >= 1; i--){
        swap(v[0], v[i]);   // 将当前最大的放置到数组末尾
        adjust(v, i, 0);  // 将未完成排序的部分继续进行堆排序
    }
}


int main()
{
    vector<int> v = {8, 1, 14, 3, 21, 5, 7, 10};
    heapSort(v, v.size());
    for(int i = 0; i < v.size(); i++){
        cout<<v[i]<<" ";
    }
    cout<<endl;

    return 0;
}

我们现在是要求最小的前K个数，那么可以先创建一个节点数为K的大根堆（0—k-1），那么从k（k—n-1）开始就跟大根堆的根结点比较，若比根节点小，则替换根节点，由上至下调整大根堆。遍历调整完最后一个根节点后，整个大根堆就是最小的前K个数。
topK代码：

class Solution {
public:
    void adjust(vector<int>& max_heap, int len, int index){
        int left = 2*index+1;
        int right = 2*index+2;

        int maxIndex = index;
        if(left < len && max_heap[left] > max_heap[maxIndex]) maxIndex = left;
        if(right < len && max_heap[right] > max_heap[maxIndex]) maxIndex = right;

        if(maxIndex != index){
            swap(max_heap[maxIndex], max_heap[index]);
            adjust(max_heap, len, maxIndex);
        }
    }
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(k == 0) return {};

        vector<int> max_heap = vector<int>(arr.begin(), arr.begin()+k);
        //建堆
        for(int i = max_heap.size()/2-1; i >= 0; i--){
            adjust(max_heap, max_heap.size(), i);
        }

        for(int i = k; i < arr.size(); i++){
            //若比大根堆的根节点要小，则调整大根堆
            if(arr[i] < max_heap[0]){
                //替换根节点
                max_heap[0] = arr[i];
                //调整大根堆
                adjust(max_heap, max_heap.size(), 0);
            }
        }

        return max_heap;
    }k
};

因为初始化堆的复杂度为O(log1) + O(log2) + O(log3) + … O(logk) = O(k)，而排序重建堆的时间复杂度为nlog(k)，所以总的时间复杂度为O(n+nlogk)=O(nlogk)

时间复杂度：O（nlogk）

思路2

这题还可以用快排的partition函数来解决，当partition函数返回的下标值等于k-1时，说明0—k-1之间都已经排好序且递增（因为partition里，是从最左边的值作为中值）
快排变形：

class Solution {
public:
    vector<int> getLeastNumbers(vector<int> &arr, int k) {
        int n = arr.size();
        if (n == k) return arr;
        if (n < k || k <= 0 || n == 0) return vector<int>();
        int left = 0, right = n - 1;
        int index = partition(arr, left, right);
        while (index != k - 1) {
            if (index > k - 1) right = index - 1;
            else left = index + 1;
            index = partition(arr, left, right);
        }
        return vector<int>(arr.begin(), arr.begin() + k);
    }

    //快排partition函数
    int partition(vector<int> &arr, int left, int right) {
        int temp = arr[left];
        while (left < right) {
            while (left < right && arr[right] >= temp) right--;
            arr[left] = arr[right];
            while (left < right && arr[left] <= temp) left++;
            arr[right] = arr[left];
        }
        arr[left] = temp;
        return left;
    }

};

参考：
https://www.cnblogs.com/wanglei5205/p/8733524.html
https://www.cnblogs.com/walker-lee/p/5814505.html