题目难度: 中等

原题链接[1]

今天继续更新 Leetcode 的剑指 Offer（专项突击版）系列, 大家在公众号 算法精选 里回复 剑指offer2 就能看到该系列当前连载的所有文章了, 记得关注哦~

题目描述

给定整数数组 nums 和整数 k，请返回数组中第 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

提示：

• 1 <= k <= nums.length <= 10^4
• -10^4 <= nums[i] <= 10^4

题目思考

1. 可以使用什么数据结构?

解决方案

思路

• 分析题目, 一个很容易想到的思路就是对数组从大到小排序, 然后第 k 个元素即为所求
• 不过这种方法需要对所有数字进行排序, 题目只要求第 k 大的, 有没有更优的方法呢?
• 第 k 问题通常都可以尝试用堆/优先队列来解决, 这道题也不例外
• 如果我们只存最大的 k 个数字到一个最小堆中, 那么只需返回堆顶即可, 无需对所有数字排序
• 这就引出了下面的思路:
• 维护一个最小堆
• 遍历数组, 将当前数字直接加入堆中
• 加入后如果堆中元素超过了 k, 就把堆顶弹出
• 由于是最小堆, 上述操作保证了堆中元素正是当前最大的 k 个数字, 更小的都被弹出去了
• 此时堆顶就是第 k 大的元素, 直接返回堆顶即可
• 另外这道题目还可以使用类似快速排序以及自定义堆的思路, 我之前的这篇文章(剑指 Offer 40. 最小的 k 个数)就包含了解决这类问题的 4 种方案和代码细节, 只是那道题稍微不同, 是求最小 k, 需要使用最大堆
• 大家感兴趣的话可以尝试参考那些思路来解决这道题, 这样举一反三可能会有新的收获~
• 下面代码中有详细的注释, 方便大家理解

复杂度

• 时间复杂度 O(NlogK): 遍历整个数组是 O(N), 每次添加数字都是操作最多 K 个元素的最小堆, 这部分是 O(logK), 所以整体就是 O(NlogK)
• 空间复杂度 O(K): 最小堆存储最多 K 个元素

代码

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 最小堆
        pq = []
        for x in nums:
            # 将当前值加入最小堆
            heapq.heappush(pq, x)
            if len(pq) > k:
                # 最小堆的元素数目超过了k, 弹出堆顶最小值
                heapq.heappop(pq)
        # 此时堆中存储的就是整个数组的最大k个数, 而堆顶就是第k大的
        return pq[0]

[1]

原题链接: https://leetcode.cn/problems/xx4gT2/