思路
方法一:排序
对原数组从小到大排序后取出前 k 个数即可。
时间复杂度:O(nlogn),其中 n 是数组 arr 的长度。算法的时间复杂度即排序的时间复杂度。
方法二:堆
我们用一个大根堆实时维护数组的前 kk 小值。首先将前 kk 个数插入大根堆中,随后从第 k+1k+1 个数开始遍历,如果当前遍历到的数比大根堆的堆顶的数要小,就把堆顶的数弹出,再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。
时间复杂度:O(nlogk),其中 n 是数组 arr 的长度。由于大根堆实时维护前 k 小值,所以插入删除都是 O(logk) 的时间复杂度,最坏情况下数组里 n 个数都会插入,所以一共需要O(nlogk) 的时间复杂度。
空间复杂度:O(k),因为大根堆里最多 k 个数。
1 class Solution {
2 public:
3 vector<int> getLeastNumbers(vector<int>& arr, int k) {
4 vector<int> res;
5 if(k <= 0)
6 return res;
7
8 priority_queue<int, vector<int>, less<int>> Q;
9
10 //先把前k个数放入大顶堆
11 for(int i = 0; i < k; ++i)
12 Q.push(arr[i]);
13
14 for(int i = k; i < arr.size(); ++i) {
15 if(arr[i] < Q.top()) {
16 Q.pop();
17 Q.push(arr[i]);
18 }
19 }
20
21 while(!Q.empty()) {
22 res.push_back(Q.top());
23 Q.pop();
24 }
25
26 return res;
27 }
28 };
方法三:利用快排的思想进行划分
这里可以参考:《算法导论(第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(n2)。
可以改进:改进选取主元的方法,使每次选出的主元在划分之后都能接近数组的中间位置,这样每次划分都能减少当前区间一半元素的工作量,可以使最坏情况下的时间复杂度降为O(n)。关于这种改进后的主元选取方法,见《算法导论(第3版)》的9.3节和这篇BFPRT算法的文章。
空间复杂度:期望为 O(logn),递归调用的期望深度为 O(logn),每层需要的空间为 O(1),只有常数个变量。最坏情况下的空间复杂度为 O(n)。最坏情况下需要划分 n次,而每层由于需要 O(1)的空间,所以一共需要 O(n) 的空间复杂度。
1 class Solution {
2 public:
3 vector<int> getLeastNumbers(vector<int>& arr, int k) {
4 if(k <= 0 || arr.empty()) return vector<int>();
5 return partition(arr, k, 0, arr.size()-1);
6 }
7
8 /*和快速排序的划分代码一模一样*/
9 vector<int> partition(vector<int>& arr, int k, int left, int right) {
10 int pivot = arr[left];
11
12 int i = left, j = right;
13 while(i < j) {
14 while(j > i && arr[j] >= pivot) --j;
15 while(i < j && arr[i] <= pivot) ++i;
16
17 swap(arr[i], arr[j]);
18 }
19
20 swap(arr[left], arr[i]);
21
22 //先判断,有选择的对左半边或者右半边进行递归
23 if(i == k-1) {
24 vector<int> res;
25 for(int t = 0; t <= i; ++t) {
26 res.push_back(arr[t]);
27 }
28 return res;
29 } else if(i > k-1) {
30 return partition(arr, k, left, i-1);
31 } else {
32 return partition(arr, k, i+1, right);
33 }
34 }
35 };