LeetCode

LeetCode 703 Kth Largest Element in a Stream 풀이: 크기 k 최소 힙

작성·코드 검증: solve together · 작성 언어: Python

0. 문제 링크와 출처

값이 하나씩 들어오는 스트림에서 매번 k번째 큰 값을 반환해야 한다. 모든 값을 매번 정렬하지 않고, 우선순위 큐의 성질을 이용해 상위 k개만 유지하면 되므로 힙 사용법을 이해해야 하지만 상태 설계가 단순해 Level 2로 판단했다.

1. 문제 요약

k와 초기 정수 배열 nums가 주어질 때, 스트림에 새로운 값이 추가될 때마다 현재까지의 k번째 큰 값을 반환하는 KthLargest 클래스를 구현한다.

중요한 조건은 다음과 같다.

  • 초기 배열은 비어 있을 수 있다.
  • 1 <= k <= nums.length + 1이다.
  • 초기 값과 새로 추가되는 값은 -10^4 이상 10^4 이하이다.
  • add는 최대 10^4번 호출된다.
  • add가 결과를 반환할 때는 스트림에 적어도 k개의 값이 존재한다.
  • 같은 값이 여러 번 들어올 수 있으며, 중복 값도 순위에 포함한다.

예를 들어 k = 3, 초기 배열이 [4, 5, 8, 2]라면 3을 추가한 뒤 정렬된 스트림은 [2, 3, 4, 5, 8]이므로 3번째 큰 값은 4다.

2. 문제 해결에 사용되는 알고리즘이나 자료구조 설명

힙 상태와 순위 판정

전체 스트림을 항상 저장하고 정렬할 필요는 없다. k번째 큰 값만 필요하므로, 지금까지 본 값 중 가장 큰 k개만 보관하면 된다.

이 k개를 최소 힙에 저장한다. 최소 힙은 가장 작은 값이 루트에 있으므로 다음 관계가 성립한다.

상위 k개 원소 중 가장 작은 원소 = 전체 원소 중 k번째 큰 원소

새 값 val을 추가하는 과정은 다음과 같다.

  1. val을 최소 힙에 넣는다.
  2. 힙의 크기가 k를 초과하면 루트, 즉 가장 작은 값을 제거한다.
  3. 현재 힙의 루트를 반환한다.

초기 배열을 처리할 때도 같은 규칙을 적용한다. 초기 배열을 먼저 모두 힙에 넣은 뒤 크기를 줄이거나, 각 값을 하나씩 추가하면서 크기를 k 이하로 유지할 수 있다.

크기 k의 최소 힙이 k번째 순위를 유지하는 방식

현재 상위 k개에 포함되지 않은 값은 아무리 작아도 k번째 큰 값에 영향을 주지 않는다. 새 값이 들어와 상위 k개에 들어오면 기존 상위 k개 중 가장 작은 값 하나가 밀려난다. 최소 힙은 바로 이 탈락 후보를 루트에서 빠르게 꺼낼 수 있다.

예를 들어 k = 3이고 힙에 [4, 5, 8]이 있다면 루트 4가 현재 3번째 큰 값이다. 10이 들어오면 [4, 5, 8, 10]에서 가장 작은 4를 제거하고 [5, 8, 10]만 남긴다. 새 루트 5가 갱신된 3번째 큰 값이다.

상태 전이와 경계 처리

초기 배열이 k개보다 적은 경우에도 생성자는 값을 그대로 힙에 저장한다. 문제 조건상 add가 값을 반환하는 시점에는 힙에 k개가 만들어져 있으므로, add에서 삽입 후 크기를 조정하면 루트가 항상 존재한다.

힙에 같은 값이 여러 개 있어도 각각 별도의 원소로 저장해야 한다. 순위는 값의 종류가 아니라 스트림에 들어온 원소의 개수로 결정되기 때문이다.

시간·공간 복잡도

초기 배열의 각 원소를 힙에 넣으며 크기를 k로 유지하므로 초기화는 O(n log k) 시간이다. add 한 번은 삽입과 최대 한 번의 삭제를 수행하므로 O(log k) 시간이다. add가 m번 호출되면 전체 갱신 시간은 O((n + m) log k)다.

힙에는 항상 최대 k개의 값만 저장하므로 공간 복잡도는 O(k)다.

3. 문제 해결 코드 (Python)

import heapq


class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.min_heap = []

        for num in nums:
            self._push(num)

    def _push(self, value: int) -> None:
        heapq.heappush(self.min_heap, value)

        if len(self.min_heap) > self.k:
            heapq.heappop(self.min_heap)

    def add(self, val: int) -> int:
        self._push(val)
        return self.min_heap[0]

_push에 삽입과 크기 제한을 모아 초기화와 add가 같은 규칙을 사용하게 했다. 힙 크기가 k보다 커지는 순간 최솟값을 제거하므로, 처리 후 힙에는 항상 현재 상위 k개만 남는다. 그래서 min_heap[0]을 별도 정렬 없이 바로 반환할 수 있다.

4. 오답 접근과 반례

모든 값을 저장하고 매번 정렬하기

새 값이 들어올 때마다 전체 스트림을 정렬하면 결과는 구할 수 있지만, 호출이 반복될수록 불필요하게 많은 값을 다시 정렬한다. 값이 n개인 시점의 정렬은 O(n log n)이고, add가 여러 번 호출되면 전체 비용이 커진다.

초기 nums = [4, 5, 8, 2], k = 3
add(3), add(5), add(10), add(9), add(4)

이런 스트림에서는 매 호출마다 전체 배열을 정렬하는 대신 상위 k개만 유지하는 것이 효율적이다.

최대 힙만 유지하기

최대 힙은 가장 큰 값을 빠르게 보여주지만, k번째 큰 값을 찾으려면 상위 원소를 여러 번 꺼내야 하거나 전체 원소를 보관해야 한다. 반대로 상위 k개를 최소 힙에 보관하면 루트가 곧 k번째 큰 값이 된다.

힙 크기를 k로 줄이지 않기

최소 힙에 모든 값을 계속 추가하면 루트는 전체 스트림의 최솟값이 된다. 이는 k가 1인 경우를 제외하면 원하는 순위와 다르다. 매 삽입 뒤 크기가 k를 초과하면 루트를 제거해야 상위 k개라는 불변식이 유지된다.

복사해 실행할 테스트

아래 검사를 풀이 코드 아래에 붙여 python 파일명.py로 실행하면 새 값 추가와 중복 값의 순위를 확인할 수 있다.

stream = KthLargest(3, [4, 5, 8, 2])
assert stream.add(3) == 4
assert stream.add(5) == 5
assert stream.add(10) == 5
assert stream.add(9) == 8
assert stream.add(4) == 8

stream = KthLargest(1, [])
assert stream.add(-5) == -5
assert stream.add(-2) == -2

5. 공식 예제와 추가 검증

작성한 KthLargest 클래스를 Python으로 실행해 공식 예제와 경계 사례를 확인했다.

공식 예제 1

입력: k = 3, nums = [4,5,8,2]
연산: add(3), add(5), add(10), add(9), add(4)
기대 출력: [4,5,5,8,8]
확인 결과: 통과
검증 목적: 새 값이 상위 k개에 들어오거나 탈락하는 과정을 확인

공식 예제 2

입력: k = 4, nums = [7,7,7,7,8,3]
연산: add(2), add(10), add(9), add(9)
기대 출력: [7,7,7,8]
확인 결과: 통과
검증 목적: 중복 값이 여러 개 있는 스트림의 순위 처리 확인

추가 검증 1: 초기 배열이 비어 있고 k가 1인 경우

입력: k = 1, nums = []
연산: add(-5), add(-2)
기대 출력: [-5,-2]
확인 결과: 통과
검증 목적: 초기 힙이 비어 있는 최소 입력 확인

추가 검증 2: 모든 값이 같은 경우

입력: k = 2, nums = [5,5]
연산: add(5), add(4)
기대 출력: [5,5]
확인 결과: 통과
검증 목적: 동일한 값도 각각 순위에 포함하는지 확인

추가 검증 3: 음수 값의 순위

입력: k = 3, nums = [-10,-5,-3]
연산: add(-7)
기대 출력: [-7]
확인 결과: 통과
검증 목적: 음수에서도 최소 힙의 상대적인 순서가 올바른지 확인

6. 다른 풀이 방식 검색 및 비교

전체 값을 정렬해 k번째 큰 값을 가져오는 방식은 구현이 직관적이지만 매 add마다 현재 스트림 전체를 처리해야 한다. 크기 k의 최소 힙은 상위 k개만 남기므로 스트림이 커져도 저장 공간을 O(k)로 제한하고, 각 추가를 O(log k)에 처리한다.

이 문제에서는 최대값 하나가 아니라 k개의 상위 원소가 필요하다. 그래서 최소 힙의 루트를 현재 k번째 큰 값으로 해석하는 것이 핵심이며, 중복 값도 별도 제거하지 않고 힙에 그대로 넣어야 한다.

마무리

Kth Largest Element in a Stream은 전체 스트림을 정렬하지 않고 순위를 유지하는 문제다. 현재 상위 k개만 최소 힙에 저장하면 힙의 최솟값이 곧 k번째 큰 값이 된다. 새 값마다 힙 크기를 k로 맞추는 불변식을 지키면 O(log k) 갱신과 O(k) 공간으로 스트림 질의를 처리할 수 있다.

#힙#LeetCode#난이도 2