Kth Largest Element in a Stream
Problem
Design a class to find the kth largest element in a stream of integers, supporting adding new elements and retrieving the kth largest element efficiently.
- 1 ≤ k ≤ 10⁴
- 0 ≤ nums.length ≤ 10⁴
- −10⁴ ≤ nums[i] ≤ 10⁴
- −10⁴ ≤ val ≤ 10⁴
- At most 10⁴ calls will be made to add.
- It is guaranteed that there will be at least k elements in the stream when you search for the kth largest element.
Example
k = 3, nums = [4,5,8,2], add(3), add(5), add(10), add(9), add(4)[null, 4, 5, 5, 8, 8]Initially, the stream is [4,5,8,2]. The 3rd largest element is 4. After add(3), stream is [4,5,8,2,3], 3rd largest is 4. After add(5), stream is [4,5,8,2,3,5], 3rd largest is 5. After add(10), stream is [4,5,8,2,3,5,10], 3rd largest is 5. After add(9), stream is [4,5,8,2,3,5,10,9], 3rd largest is 8. After add(4), stream is [4,5,8,2,3,5,10,9,4], 3rd largest is 8.
Approach
Straightforward Solution
A naive approach is to keep all elements in a list and sort it after each insertion, resulting in O(n log n) time per add call, which is inefficient for large streams.
Core Observation
The kth largest element in a stream is the element that would be at index len(stream) - k if the stream were sorted. Maintaining a full sorted list after each insertion is inefficient. Instead, a min-heap of size k can track the top k largest elements, with the smallest among them at the root representing the kth largest overall.
Path to Optimal
PreviewThe key recognition signals are 'kth largest element', 'stream', and 'add elements dynamically'. These indicate a heap or priority queue pattern because a heap can efficiently maintain the top k elements with O(log k) insertion and removal…
Full step-by-step walkthrough on Pro →
Optimal Approach
PreviewInitialize a min-heap with the first k elements from nums (or fewer if nums has less than k elements), popping the smallest elements until the heap size is k. For each add(val), push val into the heap and pop the smallest element if the heap exceeds size k…
Full step-by-step walkthrough on Pro →
Want the full reasoning chain?
Unlock the complete walkthrough, line-by-line analysis, and recall drill.
Unlock ProTime
O(n log k) for initialization and O(log k) per add call
Heapifying the initial list takes O(n), but popping elements to maintain size k costs O((n-k) log n). Each add operation involves a push and possibly a pop on a heap of size k, each O(log k).
Space
O(k)
The min-heap stores at most k elements at any time, so auxiliary space scales linearly with k, independent of the total number of elements added.
Pattern Spotlight
Heap / Priority Queue (Fixed-Size Min-Heap for Top K Elements)
Maintain a min-heap of size k to track the top k largest elements; the heap root always represents the kth largest, enabling efficient updates and queries in a dynamic stream.
Solution
| 1 | import heapq
|
| 2 |
|
| 3 | class KthLargest:
|
| 4 | def __init__(self, k: int, nums: list[int]):
|
| 5 | self.min_heap = nums
|
| 6 | self.k = k
|
| 7 | heapq.heapify(self.min_heap)
|
| 8 | while len(self.min_heap) > k:
|
| 9 | heapq.heappop(self.min_heap)
|
| 10 |
|
| 11 | def add(self, val: int) -> int:
|
| 12 | heapq.heappush(self.min_heap, val)
|
| 13 | if len(self.min_heap) > self.k:
|
| 14 | heapq.heappop(self.min_heap)
|
| 15 | return self.min_heap[0] |
Step-by-Step Solution
Build and Maintain a Min-Heap of Size k from Initial Stream
| 5 | self.min_heap = nums
|
| 6 | self.k = k
|
| 7 | heapq.heapify(self.min_heap)
|
| 8 | while len(self.min_heap) > k:
|
| 9 | heapq.heappop(self.min_heap)
|
Objective
To initialize the min-heap with the k largest elements from the initial list, ensuring efficient retrieval of the kth largest element.
Key Insight
Heapifying the initial list transforms it into a valid min-heap in O(n) time. Removing elements until the heap size is k discards smaller elements, leaving only the top k largest. This setup guarantees that the heap root is the kth largest element before any additions.
Interview Quick-Check
Core Logic
Heapify the input list to create a min-heap, then pop the smallest elements until only k remain, ensuring the heap contains the k largest elements.
State & Boundaries
Maintain the heap size at most k to keep the kth largest element at the root.
Common Pitfalls & Bugs
Failing to reduce the heap size to k after heapify can cause incorrect kth largest results.
Add New Elements and Adjust Heap to Maintain kth Largest
To insert new elements into the min-heap and remove the smallest if the heap exceeds size k, preserving the kth largest element at the root.
1 more step with full analysis available on Pro.
Line Analysis
This solution has 4 Critical lines interviewers watch for.
heapq.heapify(self.min_heap)Transform the list min_heap into a valid min-heap in-place.
Heapifying the list in O(n) time prepares it for efficient insertion and removal operations, enabling the heap to maintain the top k elements dynamically.
while len(self.min_heap) > k:While the heap size exceeds k, remove the smallest element.
This loop ensures the heap only contains the k largest elements by discarding smaller ones, which is essential for the heap root to represent the kth largest element.
if len(self.min_heap) > self.k:If the heap size exceeds k, remove the smallest element.
This conditional removal keeps the heap size fixed at k, ensuring the root remains the kth largest element after the addition.
Full line-by-line criticality + rationale for all 9 lines available on Pro.
Test Your Understanding
Why does maintaining a min-heap of size k guarantee that the root is the kth largest element in the stream?
See the answer with Pro.
Related Problems
Heap / Priority Queue pattern
Don't just read it. Drill it.
Reconstruct Kth Largest Element in a Stream from memory until it sticks. AlgoDrill blanks out key lines and makes you fill them back in, step by step.
Unlock the Kth Largest Element in a Stream drill