剑指 Offer 40. 最小的k个数

思路#

方法一：排序#

对原数组从小到大排序后取出前 k 个数即可。

时间复杂度：O(nlogn)，其中 n 是数组 arr 的长度。算法的时间复杂度即排序的时间复杂度。

 

方法二：堆#

我们用一个大根堆实时维护数组的前 kk 小值。首先将前 kk 个数插入大根堆中，随后从第 k+1k+1 个数开始遍历，如果当前遍历到的数比大根堆的堆顶的数要小，就把堆顶的数弹出，再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。

时间复杂度：O(nlogk)，其中 n 是数组 arr 的长度。由于大根堆实时维护前 k 小值，所以插入删除都是 O(logk) 的时间复杂度，最坏情况下数组里 n 个数都会插入，所以一共需要O(nlogk) 的时间复杂度。

空间复杂度：O(k)，因为大根堆里最多 k 个数。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> res;
        if(k <= 0)
            return res;
            
        priority_queue<int, vector<int>, less<int>> Q;
        
        //先把前k个数放入大顶堆
        for(int i = 0; i < k; ++i)
            Q.push(arr[i]);
        
        for(int i = k; i < arr.size(); ++i) {
            if(arr[i] < Q.top()) {
                Q.pop();
                Q.push(arr[i]);
            }
        }

        while(!Q.empty()) {
            res.push_back(Q.top());
            Q.pop();
        }

        return res;
    }
};

 

方法三：利用快排的思想进行划分#

这里可以参考：《算法导论(第3版)》的9.2和9.3节。

我们可以借鉴快速排序的思想。我们知道快排的划分函数每次执行完后都能将数组分成两个部分，小于等于分界值 pivot 的元素的都会被放到数组的左边，大于pivot的都会被放到数组的右边。

与快速排序不同的是，快速排序会根据分界值的下标递归处理划分的两侧，而这里我们只处理划分的一边，并返回分界值pivot的下标i。但是划分代码和快速排序中的代码是一模一样的。

所以算法步骤如下：

（1）编写partition()函数，与快速排序的元素划分代码一模一样，选择区间最左侧的元素为pivot，将一个区间内所有小于pivot的数放在左边，大于pivot的数放在右边，然后返回pivot的下标i。

（2）对每次划分之后pivot的下标和k进行判断：

如果划分之后的下标pivotIndex刚好等于要找的第k个数的下标(下标从零开始，所以此处为k-1), 说明arr[pivotIndex]就是要找的第k个数，又由于pivotIndex左边的数都小于arr[pivotIndex]，

所以左边的数就是最小的k个数，则将这k个数装入答案数组；

 

如果划分之后的下标pivotIndex大于要找的第k个数的下标，则递归处理左半区间[left, pivotIndex-1] ；

如果划分之后的下标pivotIndex小于要找的第k个数的下标，则递归处理右半区间[pivotIndex+1, right] 。

 

时间复杂度：期望为 O(n)，具体证明可以参考《算法导论 (第3版)》第 9.2节。

对于最好的情况：每次所选的pivot划分之后正好在数组的正中间，那么递归方程为T(n) = T(n/2) + n，解得T(n) = O(n)，所以此时此算法是O(n)线性复杂度的。

对于最坏的情况：每次的划分点都是最大值或最小值，即每次所选的pivot划分之后都好在数组最边上，一共需要划分 n - 1次，而每次划分需要O(n)的时间复杂度，所以此时此算法时间复杂度为O(n²)。

可以改进：改进选取主元的方法，使每次选出的主元在划分之后都能接近数组的中间位置，这样每次划分都能减少当前区间一半元素的工作量，可以使最坏情况下的时间复杂度降为O(n)。关于这种改进后的主元选取方法，见《算法导论(第3版)》的9.3节和这篇BFPRT算法的文章。

空间复杂度：期望为 O(logn)，递归调用的期望深度为 O(logn)，每层需要的空间为 O(1)，只有常数个变量。最坏情况下的空间复杂度为 O(n)。最坏情况下需要划分 n次，而每层由于需要 O(1)的空间，所以一共需要 O(n) 的空间复杂度。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(k <= 0 || arr.empty())  return vector<int>();
        return partition(arr, k, 0, arr.size()-1);
    }

    /*和快速排序的划分代码一模一样*/
    vector<int> partition(vector<int>& arr, int k, int left, int right) {
        int pivot = arr[left];

        int i = left, j = right;
        while(i < j) {
            while(j > i && arr[j] >= pivot) --j;   
            while(i < j && arr[i] <= pivot) ++i;

            swap(arr[i], arr[j]);
        }

        swap(arr[left], arr[i]);

        //先判断，有选择的对左半边或者右半边进行递归
        if(i == k-1) {
            vector<int> res;
            for(int t = 0; t <= i; ++t) {
                res.push_back(arr[t]);
            }
            return res;
        } else if(i > k-1) {
            return partition(arr, k, left, i-1);
        } else {
            return partition(arr, k, i+1, right);
        }
    }
};