[Python CT#3] K번째 큰 요소

 

 

#3 K번째 큰 요소

 

 

Q.  숫자로 구성된 배열 nums와 integer k가 주어집니다. nums 배열에서 k번째로 큰 요소를 찾아주세요

(sort을 사용하지 않고 구현해주세요)

 

 

 

 

 

 

정답 코드

 

import heapq

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


nums = [1, 1, 1, 2, 2, 3]
k = 2

sol = Solution()
result = sol.findKthLargest(nums, k)
print(result) # 2  nums배열에서 2번째(k)로 큰 요소 반환!

 

import heapq

 

풀이 과정!

 

파이썬은 heapq 모듈을 통해서 힙구조의 데이터 편집에 용이합니다.

 

다양한 매서드들이 있지만 오늘 그중에 사용할 매서드는 

 

heapq.nlargest() 입니다. 

 

공식문서의 설명 먼저 보겠습니다.

heapq.nlargest(n, iterable, key=None)
 iterable에 의해 정의된 데이터 집합에서 n 개의 가장 큰 요소로 구성된 리스트를 반환합니다.
key가 제공되면 iterable의 각 요소에서 비교 키를 추출하는 데 사용되는 단일 인자 함수를 지정합니다
(예를 들어, key=str.lower). 다음과 동등합니다: sorted(iterable, key=key, reverse=True)[:n].

 

heapq.nlargest(nums, k)는 nums 배열안에 있는 요소들을 최소힙(Max min)의 데이터 구조로 만들어 줍니다.

 

두번째 인자인 k는 결과값이 담긴 배열에 몇개의 요소를 담을지 결정하는 k입니다.

 

여기서 아리송하실 수 있는데, 아래 코드를 통해 이해해봅시다!

 

nlargest()의 배열 반환!

 

 

위 코드를 보시면 heapq.nlargest(k,nums)에서 [-1]을 제외시키고 k의 값을 3으로 지정합니다.

 

result를 출력시켜보면 보시는대로 [3, 2, 2] 배열이 출력됩니다.

 

다시말해 heapq.nlargest(k,nums)은 nums를 최소힙(Min heap)으로 정돈시키고 k개의 배열로 새로 만들어 주는 것입니다.

 

원래대로 heapq.nlargest(k,nums)[-1]를 적용시키면 [3,2,2] 배열에서 끝자리인 2가 반환되는 것이죠!

 

k번째로 큰 요소를 반환하는 로직이지만 사실 작은 순서로 구성된 힙에 k개의 요소를 배열로 만들고 마지막 요소를 데려오면

 

저희가 원하는 k번째로 큰 수를 찾아올 수 있는 것입니다.

 

 

# return

 

return heapq.nlargest(nums, k)[-1]