【剑指offer】最小的K个数的几种解法

题目：

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

常见思路：

这道题最简单的思路就是先把输入的n个整数排序，排序之后位于最前面的k个数就是最小的k个数。常见的排序算法都可以使用，时间复杂度就是排序的时间复杂度，较好的时间复杂度为$O(nlogn)$，这里顺便提一下python中的内置sort函数，使用的是蒂姆排序，结合了归并排序和插入排序，最坏的时间复杂度为$O(nlogn)$。
下面介绍几种时间复杂度小于$O(nlogn)$的算法。

1.$O(nlogk)$的算法：

这种思想主要就是维护一个大小为k的数据容器；首先创建一个大小为k的数据容器来存储最小的k个数字，接下来我们每次从输入的n个整数中读取一个数，比较待插入的整数和容器中的最大值，如果比已有的最大值小，就插入容器替换这个最大值；否则就不进行操作，因为这个数比容器的最大值还要大。
我们需要对容器做的操作有：

1. 在k个整数中找到最大数。
2. 在容器中删除最大值。
3. 插入一个新的数字。

如果用一个二叉树来实现这个数据容器，那么我们可以在$O(logk)$时间内实现这三个步骤。因此对于n个输入数字而言，总的时间效率就是$O(nlogk)$，这种思路特别适合处理海量数据。

这个数据容器有多种实现方式，下面介绍主要的几种实现方式。

1.1 最大堆：

我们可以选择不同的二叉树来实现这个数据容器。在最大堆中，根结点的值总是大于它的子树中任意结点的值。于是我们可以在$O(1)$时间得到已有的k个数字中的最大值，但需要$O(logk)$时间完成删除及插入操作。
下面是实现的代码：通过数组模拟最大堆。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        if(input.empty() || input.size() < k || k <= 0) return res;
        for(int i = 0; i < k; ++i){
            res.push_back(input[i]);
        }
        for(int i = k/2-1; i >= 0; i--){ // 初始化堆
            adjustHeap(res, i, k);
        }
        for(int i = k; i < input.size(); i++){
            if(input[i] < res[0]){   // 存在更小的数字时
                res[0] = input[i];
                adjustHeap(res, 0, k);   // 重新调整堆
            }
        }
        return res;
    }
    void adjustHeap(vector<int> &input, int i, int length){//调整堆
        int temp = input[i], j;
        for(j = 2*i + 1; j < length; j = j*2 + 1){
            // 沿关键字较大的孩子结点向下筛选
            if(j < length && input[j] < input[j+1]){
                ++j;   // 较大关键字的下标
            }
            if(temp >= input[j])
                break;
            input[i] = input[j];
            i = j;
        }
        input[i] = temp;
    }
    //堆排序
    void Heapsort(vector<int> &input, int length){
        for(int i = length/2 - 1; i >= 0; i--){
            adjustHeap(input, i, length);    //初始化堆
        }
        for(int i = length-1; i >= 0; i--){
            swap(input[i], input[0]);
            adjustHeap(input, 0, i);    // 重新调整堆
        }
    }
};

也可以直接使用STL实现中的堆排序，STL中并没有实现堆这种数据结构，但是algorithm中实现了堆排序算法。主要就是四个函数，make_heap()，pop_heap()，push_heap()，sort_heap()。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        if(input.empty() || input.size() < k || k <= 0) return res;
        // 初始化
        for(int i = 0; i < k; ++i){
            res.push_back(input[i]);
        }
        // 建堆
        make_heap(res.begin(), res.end());
        for(int i = k; i < input.size(); ++i){
            if(input[i] < res[0]){
                // 出堆，然后再删除
                pop_heap(res.begin(), res.end());
                res.pop_back();
                res.push_back(input[i]);
                push_heap(res.begin(), res.end());
            }
        }
        // 堆排序
        sort_heap(res.begin(), res.end());
        return res;
    }
};

1.2 优先队列：

STL中实现了优先队列，其中优先队列就是基于堆实现的。priority_queue是带权值的queue，和常规的queue类似，只能在一端入队（push），一端出队（pop）,不同的是每次元素入队之后，在容器内部按照一定次序排列，使得每次出队的元素始终是当前权值的极大值。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        priority_queue<int> pq;   //定义优先队列
        vector<int> res;
        if(input.empty() || input.size() < k || k <= 0) return res;
        for(int i = 0; i < input.size(); ++i){
            if(pq.size() < k){
                pq.push(input[i]);
            }else if(input[i] < pq.top()){
                pq.pop();   // 将元素出队
                pq.push(input[i]);  // 入队
            }
        }
        // 取出优先队列中的元素
        while(!pq.empty()){
            res.push_back(pq.top());
            pq.pop();
        }
        return res;
    }
};

1.3 红黑树：

我们还可以采用红黑树来实现我们的容器，在STL中set和multiset都是基于红黑树实现的。红黑树通过把节点分为红、黑两种颜色并根据一些规则确保树在一定程度上是平衡的，从而保证在红黑树中查找、删除和插入操作都只需要$O(logk)$时间。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        int length = input.size();
        if(input.empty() || k > length || k <=0) return res;

        multiset<int, greater<int>> leastNumbers;   // 从大到小排序
        multiset<int, greater<int>>::iterator setInterator;  // 迭代器

        vector<int>::iterator iter = input.begin();
        leastNumbers.clear();  // 初始化
        for(; iter != input.end(); ++iter){
            // 前k个数直接插入
            if(leastNumbers.size() < k){
                leastNumbers.insert(*iter);
            }else{
                setInterator = leastNumbers.begin();
                // 比较堆顶元素和要插入的元素之间的关系
                if((*leastNumbers.begin()) > (*iter)){
                    leastNumbers.erase(setInterator); // 删除堆顶元素
                    leastNumbers.insert(*iter);
                }
            }
        }
        for(setInterator = leastNumbers.begin(); setInterator != leastNumbers.end(); ++setInterator){
            res.push_back(*setInterator);
        }
        return res;
    }
};

2.$O(n)$的算法：

基于快速排序的划分思想，使用partition对数组进行划分；如果基于数组的第k个数字来划分，使得比第k个数字小的所有数字都位于数组的左边，比k个数字大的所有数字都位于数组的右边。这样调整之后，位于数组中的左边的k个数字就是最小的k个数字。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> res;
        int length = input.size();
        // 特殊情况判断！！！超时的原因在这
        if(input.empty() || k > length || k<=0) return res;
        int start = 0, end = length - 1;
        int index = Partition(input, start, end); // 对数组进行划分
        // 有点类似于二分的思想，再次划分
        while(index != k - 1){
            if(index > k - 1){
                end = index - 1;
                index = Partition(input, start, end);
            }else{
                start = index + 1;
                index = Partition(input, start, end);
            }
        }
        for(int i=0; i < k; i++){
            res.push_back(input[i]);
        }
        return res;
    }

    // 需要修改传递的数组，使用引用形式(超时)
    int Partition(vector<int> &input, int start, int end){
        int pivot = input[start];
        // 将比枢轴小的数调整到数组前面，比枢轴大的数调整到数组后面
        while(start < end){
            while(start < end && pivot <= input[end])
                end--;
            input[start] = input[end];
            while(start < end && pivot >= input[start])
                start++;
            input[end] = input[start];
        }
        input[start] = pivot;
        return start;
    }
    // 划分函数的另一种写法
    int Partition(vector<int> &input, int start, int end){
        int index = start; // 枢轴
        //int index = RandInRange(start, end);  // 随机选择一个数作为枢轴
        swap(input[index], input[end]);
        int small = start - 1;
        for(index = start; index < end; ++index){
            if(input[index] <= input[end]){
                ++small;
                if(small != index)
                    // 将小于枢轴的数交换到前面
                    swap(input[small], input[index]);
            }
        }
        ++small;  
        swap(input[small], input[end]);
        return small;   // 最后返回枢轴的位置
    }
    // 生成随机的枢轴
    int RandInRange(int a, int b){
        int c;
        c = a + rand()%(b - a + 1);
        return c;
    }
    void swap(int &fir, int &sec){
        int temp;
        temp = fir;
        fir = sec;
        sec = temp;
    }
};

参考链接：

《剑指offer》
更多代码可以查看：https://github.com/whjkm/Coding-Interviews
关于堆和优先队列的详解可以查看：
https://blog.csdn.net/zhangxiao93/article/details/51330582
https://blog.csdn.net/xiajun07061225/article/details/8553808