题目
在未排序的数组中找到第 k 个最大的元素。请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。
示例 1:
输入: [3,2,1,5,6,4] 和 k = 2
输出: 5
示例 2:
输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4
说明:
你可以假设 k 总是有效的,且 1 ≤ k ≤ 数组的长度。
思路1:快速选择
可以借用快速排序中的函数 partition(vector<int>& nums, int l, int r)
,该函数可以对数组 nums
的指定区间(左闭右闭)完成一趟快排(默认从小到大排列),并返回这一趟结束后主元的下标。
假设 partition
的返回值是 index
,其实我们能得到这样的信息:经历这一次快排过后,nums[index]
是数组中第 index+1
小的数。
题目要求找到第 k
大的元素,也就是第 nums.size()-k+1
小的元素,也就是排序后下标为 nums.size()-k
的元素。
令 target = nums.size()-k
,如果一趟结束后 index
小于 target
,那么就要对主元左侧的区间再快排,反之就对主元右侧的区间快排,直到 index
等于 target
为止,我们就找到了目标。
时间复杂度:
O
(
n
)
O(n)
O(n)
空间复杂度:
O
(
l
o
g
n
)
O(log\ n)
O(log n)
C++ 代码
class Solution {
public:
int partition(vector<int>& nums, int l, int r) {
swap(nums[l], nums[l + rand() % (r - l + 1)]);
int pivot = nums[l];
while (l < r) {
while (l < r && nums[r] > pivot)
--r;
nums[l] = nums[r];
while (l < r && nums[l] <= pivot)
++l;
nums[r] = nums[l];
}
nums[l] = pivot;
return l;
}
int findKthLargest(vector<int>& nums, int k) {
srand(time(nullptr));
int l = 0, r = nums.size() - 1, mid, target = nums.size() - k;
mid = partition(nums, l, r);
while (mid != target) {
if (mid > target) {
r = mid - 1;
} else {
l = mid + 1;
}
mid = partition(nums, l, r);
}
return nums[mid];
}
};
思路2:维护大小为 k 的最小堆
使用 STL 中的 priority_queue
实现最小堆,遍历数组并维护一个大小为 k 的最小堆即可得到答案。
C++ 代码
class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
int size = nums.size();
if (size == 1)
return nums[0];
priority_queue<int, vector<int>, greater<int> > q;
for (int i = 0; i < size; ++i) {
if (q.size() < k) {
q.push(nums[i]);
} else if (nums[i] > q.top()) {
q.push(nums[i]);
q.pop();
}
}
return q.top();
}
};