📅  最后修改于: 2023-12-03 15:26:38.613000             🧑  作者: Mango
在很多算法问题中需要查找一个数组中第 k 个最大或最小的元素,我们可以使用堆排序的方法来解决这个问题。
具体思路如下:
具体实现代码如下:
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
int findKthLargest(vector<int>& nums, int k) {
priority_queue<int, vector<int>, greater<int>> pq; // 小根堆
for(int i = 0; i < k; i++) {
pq.push(nums[i]);
}
for(int i = k; i < nums.size(); i++) {
if(nums[i] > pq.top()) {
pq.pop();
pq.push(nums[i]);
}
}
return pq.top();
}
int main() {
vector<int> nums = {3,2,1,5,6,4};
int k = 2;
int ans = findKthLargest(nums, k);
cout << ans << endl; // 输出 5
return 0;
}
其中 findKthLargest
函数的参数为一个整数数组 nums
和一个整数 k
,它返回数组中第 k
大的元素。
我们使用 std::priority_queue
类来实现小根堆。其构造函数参数分别为堆中元素类型,底层容器类型和比较函数对象类型。我们这里使用了 std::greater<int>()
的比较函数对象来实现小根堆。我们也可以使用 std::less<int>()
来实现大根堆。
从代码实现上可以看出,在已经维护好大小为 K 的小根堆的基础上,我们只需对新增的元素进行判断,如果比堆中最小的元素大,则取代堆中最小的元素,并进行小根堆的调整。能否成为小根堆中的一员自然是以它与当前堆中最小值进行比对了。由此可知小根堆的大小不变,始终为 K。
最后,返回小根堆的堆顶即可。