[리트코드] 215. 배열의 K번째 큰 요소 (Python)

Kth Largest Element in an Array - LeetCode

문제

배열과 숫자를 입력받고 배열 중 입력받은 K번째 큰 요소를 추출해라

 

과정

1. heapq 모듈을 이용해서 문제 풀이
2. heapq에서는 최대 힙을 제공하지 않는다. 따라서 부호를 변경하는 방법을 사용해서 최대 힙을 구현한다.
3. 부호를 바꿔서 최소 힙에 넣어준 뒤에 최솟값부터 pop을 해줄 때 다시 부호를 바꿔주면 최대 힙과 동일하다.
4. heapq 모델은 최소힙을 제공하므로 음수로 바꿔 heappop에서 가장 큰 값이 먼저 빠지도록 구현한다. 
5. 답을 내놓을 때는 원래 값으로 return해야 하므로, -heapq.heappop(heap)으로 부호를 원상복구한다.

입력

[3,2,1,5,6,4]

출력

2

 

풀이

import heapq
from typing import List

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)

 

한 줄씩 코드 해석해보기

import heapq
from typing import List

• heapq 모듈 import
• 타입 어노테이션을 추가하면 원소에 대한 타입까지 지정할 수 있다.

heap = list()
for n in nums:
    heapq.heappush(heap, -n)

• heap 리스트를 만들어준다.
• heapq 모델은 최소힙을 제공하므로 넣어주는 요소의 값을 음수로 바꿔서 이용한다.
• heappush : 힙 불변성을 유지하면서 -n값을 heap으로 push 해주는 메소드이다.
• 위 코드는 heapq.heapify(nums)로 대체 할 수 있다.
• heapq.heapify(nums) : 기존 리스트를 힙으로 변환해주는 메서드이다.

for _ in range(1, k):
    heapq.heappop(heap)

• 입력받은 k번째 전까지 heappop을 통해 요소를 제거한다.
• heappop : 힙 불변성을 유지하면서, heap에서 가장 작은 항목을 pop하고 반환해주는 메소드이다.

return -heapq.heappop(heap)

• heappop을 통해 요소를 반환한다.
• 답을 내놓을 때는 원래 값으로 return해야 하므로, -heapq.heappop(heap)으로 부호를 원상복구한다.

 

한 줄 풀이 (heapq)

import heapq

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

• nlargest() : k번째 큰 값을 추출하는 heapq 모듈의 기능 중 하나. k번째만큼 큰 값이 가장 큰 값부터 순서대로 리스트로 리턴된다.
• nsmallest() : n번째 작은 값을 추출하는 heapq 모듈의 기능 중 하나

 

한 줄 풀이 (sorted)

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

• 정렬을 이용해 배열을 정렬 후 뒤집은 다음 k-1 인덱스를 추출