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

Kth Largest Element in an Array

• https://leetcode.com/explore/interview/card/top-interview-questions-medium/110/sorting-and-searching/800/
• 분류: Top Interview Questions - Medium - Sorting and Searching

Explore - LeetCode

문제

• 주어진 정수 배열 nums와 정수 k에 대해 배열에서 k번째로 큰 원소를 리턴하라.
  • 정렬된 순서에 따른 k번째 원소이지 k번째 distinct element가 아님
• O(n) 시간 복잡도로 해결해야 함
• type nums: List[int] / type k: int / rtype: int

예시

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

조건

• 1 <= k <= nums.length <= 105
• -10^4 <= nums[i] <= 10^4

풀이

1. sort() 메소드 사용

내림차순으로 기본 sort 후 nums[k-1] 반환 → 시간 복잡도 O(nlogn)이므로 문제 조건에 맞지 않음

class Solution(object):
    def findKthLargest(self, nums, k):
        nums.sort(reverse=True)
        return nums[k-1]

Time: O(nlogn) / Space: O(1)

Runtime: 418 ms (beats 84.02%) / Memory Usage: 23.4 MB (beats 74.55%)

 

2. 최대 힙 (우선순위 큐) 이용

주어진 배열을 최대 힙에 저장하고 k번째 원소 출력

(1) 최대 힙에 저장하기 위해 값을 (-num, num)의 튜플로 저장

- 최소 힙을 최대 힙으로 이용하기 위해 우선순위에 -num을 둔 것

- 값을 가져올 때는 heapq.heappop(max_heap)[1]로 num(원래 값)에 접근

class Solution(object):
    def findKthLargest(self, nums, k):
        max_heap = []
        for item in nums:
            heapq.heappush(max_heap, (-item, item))
        for _ in range(k):
            ans = heapq.heappop(max_heap)[1]
        return ans

Time: O(nlogn) / Space: O(n)

Runtime: 677 ms (beats 45.15%) / Memory Usage: 31.7 MB (beats 9.17%)

 

(2) 최대 힙에 저장하기 위해 nums의 모든 원소의 부호를 뒤집은 후 heapify로 nums를 힙으로 만듦

- 값을 pop할 때 다시 - 부호를 붙여 원래 값으로 가져옴

class Solution(object):
    def findKthLargest(self, nums, k):
        for i in range(len(nums)):
            nums[i] *= -1
        heapq.heapify(nums)
        for i in range(k):
            ans = -heapq.heappop(nums)
        return ans

Time: O(n) vs O(klogn) / Space: O(1) - 기존 nums를 힙으로 바꾸고 ans 변수만 이용

Runtime: 766 ms (beats 41.06%) Memory Usage: 22.8 MB (beats 96.36%)

 

for문 대신 map 함수로 부호를 바꿨을 때

class Solution(object):
    def findKthLargest(self, nums, k):
        nums[:] = list(map(lambda x: x*-1, nums))
        heapq.heapify(nums)
        for i in range(k):
            ans = -heapq.heappop(nums)
        return ans

Runtime: 752 ms (beats 42.35%) / Memory Usage: 22.7 MB (beats 98.11%)

• 기존 nums 배열에 값을 복사해주는 식으로 해야 새로운 배열을 할당하지 않아서 메모리가 낭비되지 않음

 

3. 최소 힙 이용 (우선순위 큐)

힙의 크기를 k로 유지

→ 힙의 크기를 유지하면서 모든 원소를 힙을 거치게 해서 힙에 상위 k개의 원소만 남게 함

    (원소를 pop할 때 힙에 있는 원소 중 가장 작은 원소가 삭제되므로)

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

3 2 1 → 3 2

3 2 5 → 3 5

3 5 6 → 5 6

4 5 6 → 5 6

def fKLHeap(nums, k):
	pq = nums[:k]
	heapq.heapify(pq)
	for x in nums[k:]:
		heapq.heappush(pq, x)
		heapq.heappop(pq)
	return pq[0]

Time: O(nlogk) Space: O(k)

Runtime: 882 ms (beats 30.9%) / Memory Usage: 23.3 MB (beats 92.27%)

 

4. Quick Select

(1) pivot을 선택하고 배열의 원소를 3개의 그룹으로 나눔

: 피벗보다 작은 원소(right), 피벗과 같은 원소(mid), 피벗보다 큰 원소(left)

(2) 각 그룹에 얼마나 많은 원소가 있는지 확인

: L, M, R이 각 그룹의 원소 수라고 할 때 k ≤ L이라면 우리가 찾는 원소는 left에 있을 것.

반대로 k > L+M이라면 right를 찾아봐야 할 것(이때 k 자리에 k-L-M 넣어야 함).

둘 다 아니라면 mid 그룹을 살펴보면 됨 - mid 그룹의 원소는 모두 같은 수(pivot)로 이루어져 있으므로 바로 mid[0]을 반환하면 됨

- 찾는 원소가 left, right에 있는 경우는 left, right를 인자로 함수를 재귀호출하면 됨

class Solution:
    def findKthLargest(self, nums, k):
        if not nums: return
        pivot = random.choice(nums)
        left =  [x for x in nums if x > pivot]
        mid  =  [x for x in nums if x == pivot]
        right = [x for x in nums if x < pivot]
        
        L, M = len(left), len(mid)
        
        if k <= L:
            return self.findKthLargest(left, k)
        elif k > L + M:
            return self.findKthLargest(right, k - L - M)
        else:
            return mid[0]

Time: O(n) (최악의 경우 O(n^2)) / Space: O(n)

Runtime: 375 ms (beats 99.92%) / Memory Usage: 23.9 MB (beats 39.09%)

 

그 외 기각한 풀이

계수 정렬 (counting sort)

계수 정렬로 내림차순 정렬된 배열 만든 후 nums[k-1] 반환

- dictionary로 카운트를 하면 내장 sorted를 써서 한 번은 정렬을 해야하므로 시간 복잡도상 기각

- 리스트를 이용한 계수 정렬을 해야함 -> 데이터가 음수도 가능하므로 까다로움

→ 코드가 너무 길어져서 기각

 

max/remove 함수 이용

max/remove 함수를 이용해 제일 큰 숫자를 하나씩 제거 - O(n) 연산을 2번씩 k번 반복해야 함

→ 시간 초과로 기각