Top-K — heap 이 항상 답은 아니다

정수 배열에서 K번째 큰 값을 찾는 문제(LeetCode 215)는 top-K 유형의 표준형이다. 이 유형이 보이면 반사적으로 heap 을 꺼내 왔는데, 같은 문제를 Quick Select 로 다시 풀면서 그 반사가 절반만 맞았다는 걸 확인했다. 두 도구는 같은 답을 다른 비용으로 찾는다. 데이터가 스트리밍으로 도착하면 heap 이 사실상 유일한 선택이고, 전체 배열이 이미 메모리에 있다면 Quick Select 가 평균적으로 더 빠르다.

전체 정렬도 물론 답을 낸다. N 이 10⁵ 수준이면 O(N log N) 정렬로도 통과한다. 다만 K번째 하나를 얻으려고 N 개 전부의 순서를 확정하는 건 필요 이상의 일이다.

Heap 크기 K 유지

heap 풀이의 핵심은 max-heap 이 아니라 min-heap 이다. 크기 K 의 min-heap 에 지금까지 본 값 중 큰 것 K 개만 남기면, 루트 heap[0] 은 그 K 개 중 최솟값, 즉 현재까지의 K번째 큰 값 후보가 된다. 처음 K 개는 비교 없이 그대로 채우고, 이후부터 새 값을 루트와 비교한다. 새 값이 루트보다 크면 루트를 버리고 새 값을 넣는다. 순회가 끝나면 루트가 답이다.

max-heap 으로 전체 N 개를 쌓고 K 번 pop 하는 방법도 있다. 답은 같지만 공간이 O(N) 으로 늘어난다. 크기 K 유지 방식은 공간이 O(K) 에 머물고, K 가 N 보다 훨씬 작을수록 이득이 커진다.

flowchart TD
    A["새 값 value"] --> B{"heap 크기 < K"}
    B -->|예| C["heappush"]
    B -->|아니오| D{"value > heap[0]"}
    D -->|예| E["heapreplace"]
    D -->|아니오| F["버린다"]

import heapq


def find_kth_largest(nums, k):
    top_k = []
    for value in nums:
        if len(top_k) < k:
            heapq.heappush(top_k, value)
        elif value > top_k[0]:
            heapq.heapreplace(top_k, value)
    return top_k[0]

heapreplace 는 pop 과 push 를 한 번의 재정렬로 처리한다. heappop 후 heappush 를 따로 부르면 내부 정렬이 두 번 일어난다. 비교 부등호가 > 인 것도 의도적이다. 새 값이 루트와 같으면 교체해도 후보가 바뀌지 않으니 건너뛴다.

이 구조의 강점은 스트리밍이다. 원소가 하나 도착할 때마다 같은 판단을 반복하면 되고, N 을 미리 알 필요가 없다. 실시간 순위표나 로그 스트림처럼 원소가 계속 도착하고 끝을 모르는 상황에서도 같은 코드가 그대로 동작하며, 매 시점 top_k[0] 을 읽으면 그 시점까지의 K번째 큰 값이 나온다. 복잡도는 원소당 O(log K), 전체 O(N log K) 다.

Quick Select

Quick Select 는 발상이 다르다. K번째 큰 값은 오름차순 정렬을 가정하면 인덱스 N - K 위치의 값이다(0-indexed). [3, 2, 1, 5, 6, 4] 에서 2번째 큰 값은 정렬 후 인덱스 4 의 5 다. 전체를 정렬하지 않고 그 위치의 값만 확정할 수 있다면 정렬 비용을 아낄 수 있고, Quick Sort 의 partition 이 그 위치를 확정해 준다.

partition 은 pivot 하나를 골라 그보다 작은 값을 왼쪽으로 모으고, pivot 을 최종 정렬 위치에 놓은 뒤 그 위치를 반환한다. 반환된 위치가 목표 인덱스와 같으면 그 값이 답이다. 다르면 답이 있는 쪽 절반에서만 다시 partition 한다. Quick Sort 가 양쪽을 모두 재귀하는 것과 달리 Quick Select 는 한쪽만 들어가고, 이 차이가 평균 복잡도를 가른다. 탐색 구간이 매번 절반씩 줄어든다고 보면 N + N/2 + N/4 + … 로 2N 을 넘지 않는다. 평균 O(N) 이다.

flowchart TD
    A["partition(lo, hi) → p"] --> B{"p == target"}
    B -->|예| C["답 nums[p]"]
    B -->|"p < target"| D["오른쪽 구간만 다시 partition"]
    B -->|"p > target"| E["왼쪽 구간만 다시 partition"]
    D --> A
    E --> A

import random


def find_kth_largest(nums, k):
    target = len(nums) - k
    lo, hi = 0, len(nums) - 1
    while True:
        pivot_pos = random.randint(lo, hi)
        nums[pivot_pos], nums[hi] = nums[hi], nums[pivot_pos]
        pivot = nums[hi]
        boundary = lo
        for i in range(lo, hi):
            if nums[i] < pivot:
                nums[boundary], nums[i] = nums[i], nums[boundary]
                boundary += 1
        nums[boundary], nums[hi] = nums[hi], nums[boundary]
        if boundary == target:
            return nums[boundary]
        if boundary < target:
            lo = boundary + 1
        else:
            hi = boundary - 1

target 은 K번째 큰 값의 정렬 후 인덱스, boundary 는 pivot 미만 구역의 다음 자리다. pivot 을 무작위로 고르는 첫 줄이 최악을 막는다. 이미 정렬된 입력에서 끝 값을 pivot 으로 고정하면 구간이 매번 한 칸씩만 줄어 O(N²) 가 된다. 무작위화하면 어떤 입력이 와도 평균 O(N) 을 기대할 수 있다.

주의할 점은 in-place 라는 것이다. 입력 배열이 제자리에서 섞인다. 원본 순서를 보존해야 한다면 복사본으로 시작해야 하고, 그러면 공간 이점은 사라진다.

선택 기준

항목	heap 크기 K 유지	Quick Select
시간 평균	O(N log K)	O(N)
시간 최악	O(N log K)	O(N²), 무작위 pivot 으로 회피
공간	O(K)	O(1), in-place
스트리밍	가능	불가, 전체 배열 필요

heap 이 적합한 경우:

데이터가 스트리밍으로 도착하거나 N 을 미리 모른다
매 시점의 K번째 값을 반복해서 조회해야 한다
원본 배열을 수정하지 않고 공간을 O(K) 로 제한하고 싶다

Quick Select 가 적합한 경우:

전체 배열이 이미 메모리에 있다
K번째 값을 한 번만 찾으면 된다
배열을 제자리에서 수정해도 무방하다

라이브 코딩이라면 구현이 짧고 복잡도가 일정한 heap 으로 먼저 통과시키고, 시간이 남으면 Quick Select 의 평균 O(N) 과 무작위 pivot 을 언급하는 순서가 낫다고 본다.

K번째 큰 값이라는 문제 문장은 같아도 데이터가 주어지는 형태에 따라 도구가 갈린다. 스트리밍으로 도착하는 데이터라면 크기 K 의 heap 을 유지하고, 전체 배열이 주어져 있고 한 번만 찾는다면 partition 으로 목표 인덱스에 수렴한다. 반사적으로 heap 을 꺼내기 전에 데이터 형태를 먼저 확인하는 습관이 이번 재풀이에서 남았다.

Heap 크기 K 유지#

Quick Select#

선택 기준#

Heap 크기 K 유지

Quick Select

선택 기준