[Neetcode/Heap] Kth Largest Element in a Stream

https://neetcode.io/problems/kth-largest-integer-in-a-stream

NeetCode

 

✅ 문제 요약

KthLargest라는 클래스를 만들어야 한다. 이 클래스는 다음 두 가지 기능을 제공한다:

1. 생성자 KthLargest(k, nums)
   → 처음에 k번째로 큰 수를 추적할 수 있도록 정수 배열 nums를 받고 초기화한다.
2. add(val) 메서드
   → 새로운 값을 스트림(데이터 흐름)에 추가하고, 현재까지의 숫자들 중에서 k번째로 큰 수를 반환한다.

📌 "k번째로 큰 수"란?

숫자들을 내림차순으로 정렬했을 때, k번째에 오는 숫자를 말한다.
예: nums = [1, 2, 3, 3], k = 2라면
정렬: [3, 3, 2, 1] → 두 번째로 큰 수는 3이다.

💡 핵심 아이디어: 최소 힙(min-heap)을 이용

• 우리는 항상 k개만 유지하면서, 그 중 가장 작은 수가 "k번째로 큰 수"가 되게한다.
• 최소 힙(min-heap)은 항상 가장 작은 수를 빠르게 꺼낼 수 있는 자료구조이다.

왜 최소 힙이 필요한가?

• 우리는 k번째로 큰 수만 알면 되므로, k개의 큰 수만 유지하면 된다.
• 그 중 가장 작은 수가 "k번째로 큰 수"가 된다.

🔍 내부 동작 (단계별 설명)

초기화: KthLargest(3, [1, 2, 3, 3])

• 처음 숫자들을 최소 힙에 넣으면서 최대 3개까지만 유지한다.
• 최종적으로 힙에는 [3, 3, 3]이 남는다.
  → 이 중 가장 작은 값이 3이므로, 현재 3번째로 큰 수는 3이다.

이후 add() 호출 시 동작:

• add(3) → 힙에 3 추가 → 총 4개가 되므로 가장 작은 값 제거 → 다시 3개 유지 → 3 반환
• add(5) → 힙에 5 추가 → 4개 중 가장 작은 3 제거 → [3, 3, 5] → 3 반환
• add(6) → 힙에 6 추가 → [3, 5, 6] → 3 반환
• add(7) → 힙에 7 추가 → [5, 6, 7] → 5 반환
• add(8) → 힙에 8 추가 → [6, 7, 8] → 6 반환

---

class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.min_heap = []
        
        for num in nums:
            self.add(num)  # Use add() to maintain heap size properly

    def add(self, val: int) -> int:
        heapq.heappush(self.min_heap, val)  # Add value to heap
        if len(self.min_heap) > self.k:
            heapq.heappop(self.min_heap)    # Keep only k largest elements
        return self.min_heap[0]             # The kth largest is at the root

🧩 TL;DR:

Yes, you can solve this without using the add() function, but only if:

• The stream of numbers (including future additions) is known ahead of time, or
• You're just asked to find the k-th largest in a static list

BUT in this specific streaming problem, where values are added one by one, add() is required to:

• Maintain the stream
• Update the data structure
• Re-calculate the k-th largest value dynamically