[알고리즘] Leetcode #215 배열의 K번째 큰 요소 (Python)

---

↓↓↓ 아래는 내 리트코드 계정 ↓↓↓

leetcode.com/Jiwon_Lee/

Jiwon Lee - LeetCode Profile

 

---

문제 설명

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

 

Code

# heapq 모듈 이용
def findKthLargest(self, nums, k):
    heap = list()
    for n in nums:
        heapq.heappush(heap, -n)
        
    for _ in range(1, k):
        heapq.heappop(heap)
        
    return -heapq.heappop(heap)

heapq 모델은 최소힙을 제공하므로 음수로 바꿔 heappop에서 가장 큰 값이 먼저 빠지도록 구현한다. 답을 내놓을 때는 원래 값으로 return해야 하므로, -heapq.heappop(heap)으로 부호를 원상복구한다. 이외에도 heapify로 리스트를 힙 성질을 만족하도록 자동 배열해주는 모듈을 사용하거나, nlargest 모듈로 리스트 중 n번째로 큰 값을 자동으로 찾아주는 기능을 사용할 수 있다. 또한, 이 문제에서는 리스트 내에 값이 추가되거나 빠지는 경우 없이 주어진 값을 그대로 사용하므로, 리스트를 한 번 sort한 후 k번째로 큰 값을 찾아주는 방법도 있다.