[Heap] LeetCode: 215. Kth Largest Element in an Array

https://leetcode.com/problems/kth-largest-element-in-an-array/

 

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

 

Example 1:

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

Example 2:

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

 

Constraints:

• 1 <= k <= nums.length <= 104
• -104 <= nums[i] <= 104

 

---

▮정렬을 이용한 풀이

처음 풀어보았을 때는 정렬을 이용하여 풀었다.

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return sorted(nums)[-k]

 

 

▮힙을 이용한 풀이

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = list()
        for n in nums:
            heapq.heappush(heap, -n)
        
        for _ in range(1, k):
            heapq.heappop(heap)
        
        return -heapq.heappop(heap)

 

heapify()는 자료구조가 힙 특성을 만족하도록 바꿔주는 연산이다.

위에서처럼 모든 값을 루프를 돌면서 푸시하지 않아도 한번에 처리할 수 있다.

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heapq.heapify(nums)
        
        for _ in range(len(nums) - k):
            heapq.heappop(nums)
        
        return heapq.heappop(nums)

 

heapq에서 가장 크거나 작은 N개의 아이템을 찾을 때 nlargest()를 사용할 수 있다.

heapq.nlargest(n, iterable, key=None) 
heapq.nsmallest(n, iterable, key=None)

>>> import heapq
>>> nums = [1, 3, 6, 34, 5, 22, 67, -3, 56, -9]

>>> print(heapq.nlargest(5, nums))
[67, 56, 34, 22, 6]

>>> print(heapq.nsmallest(3, nums))
[-9, -3, 1]

 

따라서 해당 문제를 푸는 코드는 아래와 같다.

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return heapq.nlargest(k, nums)[-1]