📜  查找流中第 K 个最大元素的 C++ 程序(1)

📅  最后修改于: 2023-12-03 15:26:38.613000             🧑  作者: Mango

查找流中第 K 个最大元素的 C++ 程序

在很多算法问题中需要查找一个数组中第 k 个最大或最小的元素,我们可以使用堆排序的方法来解决这个问题。

具体思路如下:

  1. 维护一个大小为 K 的小根堆,初始化堆中元素为数组的前 K 个数。
  2. 遍历数组中剩下的元素,对于每个元素,如果它比堆中最小的元素大,则将堆中最小的元素替换为它,并进行小根堆的调整以保证堆的大小不变。
  3. 最终堆中的元素即为数组中第 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。

最后,返回小根堆的堆顶即可。