[리트코드] Kth Largest Element in an Array(python)

리트코드, K-th Largest Element in an Array_k번째 큰 원소

TL;DR

• 힙(heap)

문제 요약

1. 주어진 숫자 배열 nums와 k에 대해서 배열에서 k번째로 큰 원소를 반환한다.
2. k번째 큰 원소는 k번째 개별 원소가 아닌 정렬된 상태에서 k번째로 큰 원소를 의미한다.

입출력 형태

입출력 예시

 

- 예시 1의 경우, 내림차순으로 정렬하면 [6, 5, 4, 3, 2, 1]이 된다. 이 중 2번째로 큰 수는 5이다.

- 예시 2의 경우, 내림차순으로 정렬하면 [6, 5, 5, 4, 3, 3, 2, 2, 1]이 된다. 이 중 4번째로 큰 수는 4이다.

풀이

우선, 단순 리스트를 정렬한 후 해당하는 원소를 조회하는 방법으로 해결할 수 있다.

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

 

리스트를 내림차순으로 정렬한 뒤 k - 1번째 원소를 조회하면 원하는 결과를 얻을 수 있다.

파이썬의 리스트 성질을 사용한다면 오름차순으로 정렬한 후 -k번째 원소를 조회하여 원하는 결과를 얻을 수 있다.

다른 풀이

힙(heap)을 사용해도 문제를 해결할 수 있다.

힙이란 트리 기반의 자료구조형으로 파이썬에서는 부모 노드가 자식 노드에 비해서 작은 값을 갖게 되는 최소 힙(min heap)으로 구현이 된다. 힙의 최소 값을 루트를 반환함으로써 얻어낼 수 있는데 이를 이용하면 답을 구해낼 수 있다.

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

 

원하는 값은 가장 큰 수이기 때문에 입력받은 숫자 배열에 -를 붙여서 heapq를 사용하여 힙을 구성하였다. 이렇게 될 경우 힙의 루트는 가장 큰 수가 음수가 된 형태가 된다. k까지 루트를 반환하면 원하는 답을 얻을 수 있다. 단, 이 때, -를 다시 붙여서 원래의 수를 반환하도록 해주어야 한다.

 

이 외에도 리스트와 같은 자료를 힙 구조로 만들어주는 heapify, 힙의 가장 큰 수들을 얻어낼 수 있는 nlargest 메소드 등을 사용해서도 문제를 해결할 수 있다.