leetcode原题链接:数组中的第K个最大元素
上一篇:HOT73-柱状图中最大的矩形
下一篇:HOT75-前K个高频元素
题目描述
给定整数数组 nums
和整数 k
,请返回数组中第 k
个最大的元素。请注意,你需要找的是数组排序后的第 k
个最大的元素,而不是第 k
个不同的元素。
你必须设计并实现时间复杂度为 O(n)
的算法解决此问题。
示例 1:
输入: [3,2,1,5,6,4],
k = 2
输出: 5
示例 2:
输入: [3,2,3,1,2,4,5,5,6],
k = 4
输出: 4
提示:
1 <= k <= nums.length <= 105
-104 <= nums[i] <= 104
解题方法: 小顶堆保存最大的k个元素,堆顶是这k个元素中最小的元素,也就是第k大的元素。求最大的k个元素用小顶堆,求最小的k个元素用大顶堆。同时注意下c++的语法糖。std::less用于定义大顶堆, std::greater用于定义小顶堆。
C++代码
#include <iostream>
#include <vector>
#include <queue>
#include <functional> // std::less, std::greater
/*
* 最大的k个元素,采用小顶堆, std::greater
* 最小的k个元素,采用大顶堆, std::less
* std::priority_queue的成员函数如下:
* empty(),size(),top(),push(), emplace()[c++11], pop(), swap(c++11)
*/
class Solution {
public:
int findKthLargest(std::vector<int>& nums, int k) {
int n = nums.size();
if (n < k) {
return -1;
}
std::priority_queue<int, std::vector<int>, std::greater<int>> q;// 小顶堆
for (int i = 0; i < k; i++) { // 初始化小顶堆
q.emplace(nums[i]);
}
for (int i = k; i < n; i++) { //遍历数组,如果当前数字比堆顶元素大,则替换堆顶元素
if (!q.empty() && q.top() < nums[i]) {
q.pop();
q.emplace(nums[i]);
}
}
return q.top();
}
};