在子区间 [left, right] 中选择第 k 大的数时,完全照搬快排的划分算法:
1.选择一个枢轴(pivot),然后交换到 left 位置
int randIdx = rand() % (right - left + 1) + left; // 随机选择 pivot
swap(nums[randIdx], nums[left]);
2.使用划分算法确定 pivot 的位置。大于 pivot 的元素移到左边,小于等于 pivot 的元素移到右边。
int pivot = nums[left];
int l = left, r = right;
while(l < r)
{
//...
}
// 一轮 partition 完成
一轮 partition 完成后,pivot 的位置为 l,此时考察 l 和 left + K 的关系,[left, right] 中第 K 个位置的下标是 left + K - 1:
l = left + K - 1: pivot 刚好在 [left, right] 第 k 个位置,找到答案了
l > left + K - 1: [left, l] 中有 l - left + 1 个数字,还要在 [l + 1, right] 中找 K - (l - left + 1) 个
l < left + K - 1: 在 [left, l - 1] 中继续找第 k 大
时间复杂度 平均O(N),最坏O(N^2),但是最坏情况太难达到了,一般还是认为快速选择算法是 O(N) 的。
快速选择算法有一个优化:BFPRT算法,也叫中位数的中位数算法,它进一步优化了 pivot 的选取方法,使得最坏时间复杂度也变为
,它由Blum、Floyd、Pratt、Rivest、Tarjan提出。
快速选择算法每完成一轮 partition 将数据分成不均匀的两份后,只有其中一份对结果有影响,可以去掉一部分,数据规模就减少一部分,所以也叫减治算法。插入排序,DFS, BFS, 拓扑排序也隐含了减治的思想:每完成一轮,下一轮就可以少考虑一个数据。
代码(c++)
class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
int n = nums.size();
if(n == 1) return nums[0];
return partition(nums, k, 0, n - 1);
}
private:
int partition(vector<int>& nums, int k, int left, int right)
{
// 在 nums 的 [left .. right] 中找第 k 大
int randIdx = rand() % (right - left + 1) + left; // 随机选择 pivot
swap(nums[randIdx], nums[left]);
int pivot = nums[left];
int l = left, r = right;
while(l < r)
{
while(l < r && nums[r] <= pivot)
--r;
if(l < r)
{
nums[l] = nums[r];
++l;
}
while(l < r && nums[l] > pivot)
++l;
if(l < r)
{
nums[r] = nums[l];
--r;
}
}
nums[l] = pivot;
if(l == left + k - 1)
return nums[l];
else if(l < left + k - 1)
return partition(nums, k - (l - left + 1), l + 1, right);
else
return partition(nums, k, left, l - 1);
}
};