题目描述:给一个整数数组和一个正整数K,返回数组中第K大的元素。
思路1:堆排序(优先队列)
维护一个小顶堆,堆的大小限制为K,堆里面装的元素就是当前数组中前K大的元素。
这个思路非常简单,用STL的priority_queue
直接就解决了,不需要过多阐述。
注意:priority_queue
默认是大顶堆,也就是top()
返回最大值,需要用greater<>
改成小顶堆。
时间复杂度:\(O(NlogK)\)
int findKthLargest(vector<int>& nums, int k) {priority_queue<int, vector<int>, greater<int>> pq;int n = nums.size();for(int i=0; i<n; ++i){pq.push(nums[i]);if(pq.size() > k){pq.pop();}}return pq.top();
}
思路2:快速选择(快速排序改)
基于优先队列的方法并不是时间最优的。事实上我们可以做到时间复杂度期望为\(O(N)\)的方法,以下是参考力扣215题解的分析:
我们先来回顾快速排序,我们对数组 \(a[l⋯r]\) 做快速排序的过程是:
-
划分: 将数组 \(a[l⋯r]\) 「划分」成两个子数组 \(a[l⋯q−1]\)、\(a[q+1⋯r]\),使得 \(a[l⋯q−1]\) 中的每个元素小于等于 \(a[q]\),且 \(a[q]\) 小于等于 \(a[q+1⋯r]\) 中的每个元素。其中,计算下标 \(q\) 也是「划分」过程的一部分。
-
解决: 通过递归调用快速排序,对子数组 \(a[l⋯q−1]\) 和 \(a[q+1⋯r]\) 进行排序。
-
合并: 因为子数组都是原址排序的,所以不需要进行合并操作,\(a[l⋯r]\) 已经有序。
由此可以发现每次经过「划分」操作后,我们一定可以确定一个元素的最终位置,即 \(x\) 的最终位置为 \(q\),并且保证 \(a[l⋯q−1]\) 中的每个元素小于等于 \(a[q]\),且 \(a[q]\) 小于等于 \(a[q+1⋯r]\) 中的每个元素。所以只要某次划分的 \(q\) 为倒数第 \(k\) 个下标的时候,我们就已经找到了答案。 我们只关心这一点,至于 \(a[l⋯q−1]\) 和 \(a[q+1⋯r]\) 是否是有序的,我们不关心。
因此我们可以改进快速排序算法来解决这个问题:在分解的过程当中,我们会对子数组进行划分,如果划分得到的 \(q\) 正好就是我们需要的下标,就直接返回 \(a[q]\);否则,如果 \(q\) 比目标下标小,就递归右子区间,否则递归左子区间。这样就可以把原来递归两个区间变成只递归一个区间,提高了时间效率。这就是“快速选择”算法。
我们使用双指针法的快排,可以有效应对各种数据:
int qSelect(vector<int>& nums, int l, int r, int k){if(l == r){return nums[k];}int partition = nums[l];int lp = l - 1, rp = r + 1;while(lp < rp){do{lp++;}while(nums[lp] < partition);do{rp--;}while(nums[rp] > partition);if(lp < rp){swap(nums[lp], nums[rp]);}}if(k <= rp){return qSelect(nums, l, rp, k);}else{return qSelect(nums, rp+1, r, k);}
}
int findKthLargest(vector<int>& nums, int k) {int n = nums.size();return qSelect(nums, 0, n-1, n-k);
}