Kth Largest Element in an Array
정수 배열 nums와 정수 k가 주어집니다. 배열을 큰 값부터 작은 값 순으로 정렬했을 때 1부터 세어 k번째 위치에 있는 값, 즉 nums에서 k번째로 큰 값을 반환하세요.
같은 값도 각각 따로 셉니다. [5, 5, 1]에서 가장 큰 값은 5이고 두 번째로 큰 값도 5입니다.
함수
- numsinteger-array
- 순위를 매길 값
- kinteger
- 반환할 가장 큰 값, 가장 큰 값의 경우 1
- 반환값integer
- 중복을 포함하여 k번째로 큰 값
제약 조건
1 ≤ k ≤ nums.length ≤ 104-104 ≤ nums[i] ≤ 104- 같은 값도 서로 별개의 값으로 셉니다.
예제
- 입력
- nums = [7, 2, 9, 4, 9, 1]k = 2
- 출력
- 9
- 설명
- 값을 큰 순서부터 작은 순서로 나열하면
9, 9, 7, 4, 2, 1입니다. 9 두 개는 각각 별도로 계산되므로, 두 번째로 큰 값은7이 아니라9입니다.
- 입력
- nums = [5, -3, 8, 0, 2]k = 4
- 출력
- 0
- 설명
- 값을 큰 것부터 작은 것 순서로 나열하면
8, 5, 2, 0, -3이고, 그중 네 번째 값은0입니다.
- 입력
- nums = [6]k = 1
- 출력
- 6
- 설명
- 값이 하나뿐이고
k = 1이면, 그 값이 가장 큽니다.
제출 시 숨은 테스트 +15개
후속 질문
이제 값이 한 번에 하나씩 들어옵니다. 각 값이 들어올 때마다 지금까지 본 모든 값의 중앙값을 값당 O(log n) 시간에 구할 수 있나요?
힌트
하나씩 열어 보세요. 열 때마다 조금 더 알려 줍니다.
큰 값부터 작은 값 순으로 정렬하면 답은 알려진 위치에 있습니다. 어느 위치일까요? 그리고 그 값을 알기 위해 다른 모든 값이 필요한가요?
k번째로 큰 값은
k개의 가장 큰 값 중 가장 작은 값입니다. 지금까지 본 값 중 가장 큰k개만 유지한다면, 새 값을 그중 어떤 값과 비교해야 할까요?최대
k개의 값을 최소 힙에 유지하세요. 새 값이 더 크면 맨 위 값을 교체하고, 마지막에 맨 위에 있는 값이 답입니다. 평균 시간 복잡도O(n)을 위해 퀵 정렬처럼 무작위 피벗을 기준으로 분할하고, 인덱스n-k를 포함하는 쪽만 유지하세요.
풀이
정렬한 뒤 한 위치의 값을 읽으면 질문에 답할 수 있고, 여기서는 충분히 빠릅니다. 면접관이 확인하려는 것은 그 정렬 과정 중 얼마나 많은 부분을 건너뛸 수 있는지입니다. 모든 n개가 아니라 한 위치만 필요하기 때문입니다. 크기가 k인 최소 힙은 답이 될 수 있는 값만 유지하며, 퀵셀렉트는 퀵소트처럼 분할하되 답이 있는 쪽만 따라가므로 평균 시간 복잡도를 O(n)으로 낮춥니다.
한 위치를 정렬하고 읽기
핵심 아이디어
k번째로 큰 값은 정렬된 순서로 정의되므로, 그 순서를 만들면 됩니다. 큰 값부터 작은 값 순으로 정렬하면 [7, 2, 9, 4, 9, 1]은 [9, 9, 7, 4, 2, 1]이 되고, k번째로 큰 값은 인덱스 k-1에 있습니다. k = 2이면 인덱스 1, 즉 두 번째 9입니다. 정렬 결과에서 가장 작은 값이 앞에 온다면 대신 인덱스 n-k를 읽으면 됩니다. [1, 2, 4, 7, 9, 9]의 인덱스 4도 같은 9입니다.
중복 값은 특별히 처리할 필요가 없습니다. 정렬하면 모든 중복 값이 유지되고, 각 값은 각자의 위치를 차지합니다.
n = 10^4일 때 정렬은 약 n log n ≈ 1.3 × 10^5번의 비교를 수행하며, 모든 테스트를 통과합니다. 낭비되는 점은 한 위치만 중요할 때도 n개의 값을 모두 정렬한다는 것입니다. 다음 두 접근법은 이 작업을 덜 수행합니다.
알고리즘
- 호출자의 배열이 원래 상태를 유지하도록
nums를 복사합니다. - 복사본을 정렬합니다. 숫자 비교를 사용하세요. 일부 언어는 기본적으로 숫자를 텍스트로 비교합니다.
- 큰 값부터 정렬한 순서에서는 인덱스
k-1을, 작은 값부터 정렬한 순서에서는 인덱스n-k를 반환합니다.
def findKthLargest(nums, k):
# Largest first: the k-th largest sits at index k-1.
ordered = sorted(nums, reverse=True)
return ordered[k - 1]최솟값 힙에서 가장 큰 k개 유지하기
핵심 아이디어
k번째로 큰 값은 k개의 가장 큰 값 중 가장 작은 값입니다. 따라서 nums를 한 번 순회하면서 지금까지 본 값 중 가장 큰 k개만 최소 힙에 보관합니다. 최소 힙의 맨 위는 가장 작은 값이며, 이것이 바로 정답 후보입니다.
값 x가 들어왔을 때 힙에 값이 k개보다 적으면 추가합니다. 그렇지 않으면 x를 맨 위 값과 비교합니다. x가 더 크지 않다면, 보관해 둔 값 중 적어도 k개가 x보다 크거나 같으므로 x는 절대 정답이 될 수 없어 건너뜁니다. x가 더 크다면 맨 위 값은 가장 큰 k개에서 빠지므로, 그 값을 x로 교체합니다. k = 4인 예제 2에서는 처음 네 값인 5, -3, 8, 0이 힙을 채우고 맨 위 값은 -3입니다. 다음으로 2가 -3보다 커서 이를 교체하고, 맨 위 값은 0이 되며, 0이 정답입니다.
각 값의 힙 연산은 최대 한 번이며 비용은 O(log k)이므로 전체 시간 복잡도는 O(n log k), 메모리 복잡도는 O(k)입니다. k가 작을 때는 정렬보다 효율적이며, 스트림에서도 작동하므로 모든 값을 한꺼번에 보유할 필요가 없습니다. Python에는 heapq, Java에는 PriorityQueue, C++에는 greater와 함께 사용하는 priority_queue, Go에는 container/heap, Rust에는 Reverse와 함께 사용하는 BinaryHeap, PHP에는 SplMinHeap이 있습니다. 다른 언어의 코드는 힙을 배열로 구현합니다. 인덱스 i의 자식은 2i+1과 2i+2에 위치하며, 1부터 세는 Lua와 R에서는 2i와 2i+1에 위치합니다.
알고리즘
- 빈 최소 힙으로 시작합니다.
- 각 값
x에 대해 힙에k개보다 적은 값이 있으면 추가합니다. k개가 되면x가 최상단 값보다 클 때만 최상단 값을x로 바꿉니다.- 마지막 값까지 처리한 후 힙의 최상단 값을 반환합니다.
import heapq
def findKthLargest(nums, k):
# A min-heap of the k largest values so far; its top is the smallest of them.
heap = []
for x in nums:
if len(heap) < k:
heapq.heappush(heap, x)
elif x > heap[0]:
heapq.heapreplace(heap, x) # drop the top, add x
return heap[0]3-way 분할을 사용한 퀵셀렉트
핵심 아이디어
퀵 정렬은 피벗을 선택하고 분할합니다. 작은 값은 왼쪽에, 큰 값은 오른쪽에 둡니다. 한 번 분할하고 나면 양쪽이 아직 정렬되지 않았더라도 피벗은 최종 정렬 위치에 놓입니다. 퀵셀렉트는 이 사실을 활용합니다. 오름차순 기준으로 답은 인덱스 target = n-k에 있습니다. 분할 후 target은 피벗의 왼쪽이나 피벗 위치, 또는 오른쪽에 있으므로 한쪽만 계속 살펴보고 다른 쪽은 버립니다.
[7, 2, 9, 4, 9, 1]과 k = 2의 경우 target은 6-2 = 4입니다. 4를 기준으로 분할하면 2와 1은 인덱스 0과 1을 차지하고, 4는 인덱스 2를 차지하며, 7, 9, 9는 인덱스 3부터 5를 차지합니다. 인덱스 4는 오른쪽에 있으므로 인덱스 3부터 5만 남깁니다. 이 값들을 9를 기준으로 분할하면 7은 인덱스 3을 차지하고 두 개의 9는 인덱스 4와 5를 차지합니다. 인덱스 4에는 9가 있으므로 답은 9입니다.
3방향 분할을 사용하세요. 피벗보다 작은 값, 피벗과 같은 값, 피벗보다 큰 값 순으로 나누고, lt와 gt로 추적합니다. 같은 값으로 이루어진 블록 [lt, gt]은 정렬된 위치에 있으므로 target이 그 안에 있으면 작업이 끝납니다. 일반적인 2방향 분할에서는 10^4개의 7이 들어 있는 배열이 한 번에 값 하나씩 줄어들어 약 5 × 10^7단계가 걸립니다. 3방향 분할은 한 번의 순회로 답을 구합니다.
피벗은 무작위로 선택하세요. 절반의 확률로 피벗은 범위의 가운데 절반에 놓이고, 이 경우 범위는 최대 4분의 3으로 줄어듭니다. 따라서 기대 작업량은 n개 값을 몇 번 순회하는 정도인 O(n)입니다. 모든 피벗이 극단값이면 최악의 경우는 여전히 O(n²)이며, 첫 번째 원소처럼 고정된 선택 방식은 정렬된 입력에서 최악의 경우를 유발합니다. 코드는 복사본에서 작동하므로 O(n)의 메모리가 필요합니다. 입력을 변경해도 된다면 nums 자체를 분할하여 메모리를 O(1)로 줄일 수 있습니다.
알고리즘
nums를a에 복사하고,target = n-k,lo = 0,hi = n-1로 설정합니다.a[lo..hi]에서 임의의 피벗을 선택합니다.a[lo..hi]를 피벗보다 작은 값, 피벗과 같은 값, 피벗보다 큰 값으로 분할하고, 같은 값은a[lt..gt]에 둡니다.target < lt이면hi = lt-1로 설정하고,target > gt이면lo = gt+1로 설정합니다. 그렇지 않으면 피벗을 반환합니다.- 2단계부터 반복합니다.
import random
def findKthLargest(nums, k):
a = list(nums)
target = len(a) - k # the answer's index once a is sorted smallest first
lo, hi = 0, len(a) - 1
while True:
pivot = a[random.randint(lo, hi)]
# Three-way partition of a[lo..hi]: < pivot, then == pivot, then > pivot.
lt, i, gt = lo, lo, hi
while i <= gt:
if a[i] < pivot:
a[lt], a[i] = a[i], a[lt]
lt += 1
i += 1
elif a[i] > pivot:
a[i], a[gt] = a[gt], a[i]
gt -= 1
else:
i += 1
# Now a[lt..gt] all equal pivot, and they are in their sorted places.
if target < lt:
hi = lt - 1
elif target > gt:
lo = gt + 1
else:
return pivot
함정과 경계 사례
오답의 대부분은 중복 값을 처리하는 방법과 위치를 세는 두 가지 방식을 혼동하는 데서 나옵니다.
- 먼저 중복 값을 제거하는 경우. 이 문제에서는 모든 복사본을 셉니다.
[7, 2, 9, 4, 9, 1]에서k = 2일 때 답은9이지만, 배열을 집합으로 바꾸면7이 됩니다. - 잘못된 인덱스를 읽는 경우.
k는 1부터 세므로, 가장 큰 값부터 정렬한 순서에서는 답이 인덱스k-1에 있고 가장 작은 값부터 정렬한 순서에서는 인덱스n-k에 있습니다.n-k-1이 아닙니다. - 숫자를 텍스트로 정렬하는 경우. JavaScript와 TypeScript에서
[10, 9, 2].sort()는[10, 2, 9]를 반환합니다.(a, b) => a - b를 전달하세요. - 크기가
k인 최대 힙을 사용하는 경우. 가장 큰 값을 제거하면 가장 작은k개의 값이 남고 k번째로 작은 값을 반환합니다. - 양방향 분할이나 고정 피벗을 사용하는 Quickselect. 같은 값이 많거나 배열이 정렬되어 있으면
O(n²)의 비용이 발생하며, 대규모 테스트에는 이런 경우가 포함됩니다.
자주 묻는 질문4
배열에서 K번째로 큰 원소의 시간 복잡도는 얼마인가요?
정렬에는 O(n log n) 시간이 걸립니다. 크기가 k인 최소 힙은 O(n log k) 시간과 O(k) 메모리가 필요합니다. 무작위 피벗을 사용하는 퀵셀렉트는 평균적으로 O(n) 시간이 걸리고 최악의 경우 O(n²) 시간이 걸리지만, 무작위 피벗을 사용하면 최악의 경우가 발생할 가능성은 매우 낮습니다.
k번째로 큰 원소를 찾을 때 최대 힙이 아니라 최소 힙을 사용하는 이유는 무엇인가요?
힙은 지금까지 확인한 가장 큰 값 k개를 저장하며, 비교한 뒤 제거해야 하는 값은 그중 가장 작은 값입니다. 최소 힙은 이 값을 맨 위에 둡니다. 최대 힙은 모든 n개 값을 넣고 k-1번 꺼내는 경우에만 작동하며, 이때 O(n)의 메모리가 필요합니다.
k번째로 큰 요소를 찾을 때 힙을 사용해야 할까요, 아니면 퀵셀렉트를 사용해야 할까요?
Quickselect는 평균적으로 더 빠르며 O(n)이지만, 모든 값을 메모리에 저장하고 순서를 재배치해야 합니다. 힙은 나쁜 최악의 경우가 없고 O(n log k)이며, 값이 한 번에 하나씩 들어오고 모든 값을 저장할 수 없는 상황에서도 사용할 수 있습니다. 면접에서는 두 방법 모두 설명하고, 후속 질문에서 요구하는 방법을 코드로 작성하세요.
k번째로 큰 원소를 최악의 경우 선형 시간에 찾을 수 있을까요?
네. 중앙값의 중앙값 규칙은 값의 일정 비율을 반드시 걸러 내는 피벗을 선택하므로, 최악의 경우에도 선택 작업이 O(n)이 됩니다. 다만 실제로는 무작위 피벗보다 느립니다. 값의 범위가 -10^4부터 10^4까지로 제한되어 있다면 각 값이 몇 번 나타나는지 세고, k개의 값을 지나칠 때까지 10^4부터 차례로 내려가도 됩니다. 이 방법의 시간 복잡도는 O(n + 2 × 10^4)입니다.
비슷한 문제
같은 아이디어를 쓰는 문제입니다. 두세 개를 풀면 패턴이 몸에 익습니다.
Python
def findKthLargest(nums, k):
# 여기에 코드를 작성하세요케이스 1
케이스 2
케이스 3
입력
nums = [7, 2, 9, 4, 9, 1] k = 2
기대값
9