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

原题

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]