HOT74-数组中的第K个最大元素

JXH_1232023-08-09 13:24

题目描述

给定整数数组 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个元素用大顶堆。同时注意下c++的语法糖。std::less用于定义大顶堆, std::greater用于定义小顶堆。

C++代码

cpp 复制代码

#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 == 0 || k > n) {
            return -1;
        }
        std::priority_queue<int, std::vector<int>, std::greater<int>> pq;
        for (int i = 0; i < n; i++) {
            if (i < k) { //初始化小顶堆上的k个元素
                pq.emplace(nums[i]);
            } else if (nums[i] > pq.top()) { //当前遍历的数字比堆顶元素大
                pq.pop();//先弹出堆顶元素
                pq.emplace(nums[i]);//再压入元素
            }
        }
        return pq.top();//小顶堆的头节点就是第k大元素
    }
};