215. Kth Largest Element in an Array

[zeikar]

Problem link

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

Problem Summary

배열에서 정렬 없이 k 번째 큰 수를 찾는 문제

Solution

정렬하면 쉽지만 안하고 풀어야 한다. quick sort를 응용한 quick select 알고리즘을 쓰면 된다.

quick sort와 비슷하게 피벗을 잡고 피벗 기준으로 왼쪽은 더 큰 값들, 오른쪽에는 더 작은 값들을 몰아 넣는다. 여기서 포인트는 quick select는 정렬이 필요가 없으므로 왼쪽, 오른쪽 기준으로 한 덩어리만 체크해주면 된다.
k가 왼쪽 배열보다 작다면 왼쪽 배열 내에서 k 번째 수를 찾으면 된다.

n + n/2 + n/4 + ... 으로 줄어서 시간 복잡도는 O(n). 물론 운이 아주 나쁘다면 O(n^2)겠지만 평균은 O(n) 이다.

메모리 제한은 따로 없으므로 구현이 간단하게 배열 복사 방식으로 구현하면 좋다. (in-place로 배열 원소 교체도 가능한데 코드가 꽤 지저분하다)

Source Code

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        def quickSelect(arr, n):
            if len(arr) == 1:
                return arr[0]
            
            pivot = arr[random.randint(0, len(arr)-1)]
            left = [x for x in arr if x > pivot]
            mid = [x for x in arr if x == pivot]
            right = [x for x in arr if x < pivot]

            if n < len(left):
                return quickSelect(left, n)
            elif n < len(left) + len(mid):
                return pivot
            else:
                return quickSelect(right, n - len(left) - len(mid))
        
        return quickSelect(nums, k-1)