Leetcode: Kth largest element in an Array

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

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

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

 

소팅되지 않은 배열에서 k번째로 큰 수를 찾는 문제이다.

가장 쉽게 떠오르는 방법은 배열을 소팅해서 끝에서 k번째 element를 찾는 것이다. 길이 N짜리 배열을 소팅하므로 time complexity는 O(nlogn)이다.

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

 

조금 더 잘 풀 수 있는 방법이 없을까 고민해보았다. Priority Queue를 사용하는 방법도 있을 텐데, 배열의 모든 수를 heapify하면 결국 O(nlogn)이 필요하다. 그렇지만 우리는 배열의 모든 수를 볼 필요가 없고, 가장 큰 k개의 element만 보면 된다. 그렇기 때문에 길이 k짜리 heap을 유지한다. 이 때, max heap이 아니라 min heap을 쓰는 것이 포인트이다. heap의 크기가 k가 넘어 원소를 뽑아 버릴 때, 버려야 하는 원소는 작은 원소이다. 즉, k=4일 때, heap에 5번째 원소가 들어왔다면 우리가 버려야 하는 원소는 그 중 가장 작은 원소라는 것이다. 나머지 4개의 원소는 top k 후보군이기 때문이다. heap의 크기가 커졌을 때, 가장 작은 원소를 매번 pop해서 버려야하므로 heap에서 이것을 O(1)으로 뽑으려면 min heap이어야 한다. 

이렇게 마지막 원소까지 순회했을 때, 마지막으로 pop한 값 (pq[0])이 k번째로 큰 원소가 된다. 해당 heap은 배열에서 가장 큰 k개의 원소가 작은 순서부터 큰 순서로 들어있기 때문이다. 

이 경우, time complexity는 O(nlogk)가 된다. 

 

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        pq = []
        for num in nums:
            heapq.heappush(pq, num)
            if len(pq) > k:
                heapq.heappop(pq)
        return heapq.heappop(pq)

 

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