[leetcode] 215. Kth Largest Element in an Array @ python

最新推荐文章于 2022-07-02 13:48:51 发布

原创最新推荐文章于 2022-07-02 13:48:51 发布 · 461 阅读

0 ·

CC 4.0 BY-SA版权

本文介绍了一种使用小根堆解决寻找未排序数组中第K大元素问题的方法。通过维护长度为K的小根堆，确保堆顶元素即为所求。代码示例使用Python的heapq库实现。

原题

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

Input: [3,2,1,5,6,4] and k = 2
Output: 5
Example 2:

Input: [3,2,3,1,2,4,5,5,6] and k = 4
Output: 4
Note:
You may assume k is always valid, 1 ≤ k ≤ array’s length.

解法

小根堆. https://docs.python.org/3.0/library/heapq.html
维持一个长度为k的小根堆, 那么第一个元素就是结果. 如果nums的长度大于k, 我们需要把多出的元素剔除.

代码

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        heapq.heapify(nums)
        while len(nums) > k:
            heapq.heappop(nums)
        return nums[0]