题目链接
给定整数数组 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
解题思路
思路:堆(优先队列)
- 遍历数组元素,对于挡圈元素 num:
- 如果优先队列中的元素个数小于 k 个,则将当前元素 num 放入优先队列中。
- 如果优先队列中的元素个数大于等于 k 个,并且当前元素 num 大于优先队列的队头元素,则弹出队头元素,并将当前元素 num 插入到优先队列中。
- 遍历完,此时优先队列的队头元素就是第 k 个最大元素,将其弹出并返回即可。
这里我们借助了 Python 中的 heapq 模块实现优先队列算法,这一步也可以通过手写堆的方式实现优先队列。
解题代码
import heapq
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
res = []
for num in nums:
if len(res) < k:
heapq.heappush(res, num)
elif num > res[0]:
heapq.heappop(res)
heapq.heappush(res, num)
return heapq.heappop(res)
参考资料:datawhalechina