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
The 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. By keeping a min-heap of size k, the root always holds the kth largest element, allowing constant-time retrieval and logarithmic-time updates.
Optimal Approach
Initialize 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. The root of the heap is then the kth largest element. This approach guarantees O(log k) time per add and O(k) space.
Time
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] |
| 1 | import java.util.PriorityQueue; |
| 2 | |
| 3 | public class KthLargest { |
| 4 | private PriorityQueue<Integer> minHeap; |
| 5 | private int k; |
| 6 | |
| 7 | public KthLargest(int k, int[] nums) { |
| 8 | this.minHeap = new PriorityQueue<>(); |
| 9 | this.k = k; |
| 10 | for (int num : nums) { |
| 11 | this.minHeap.offer(num); |
| 12 | } |
| 13 | while (this.minHeap.size() > k) { |
| 14 | this.minHeap.poll(); |
| 15 | } |
| 16 | } |
| 17 | |
| 18 | public int add(int val) { |
| 19 | this.minHeap.offer(val); |
| 20 | if (this.minHeap.size() > this.k) { |
| 21 | this.minHeap.poll(); |
| 22 | } |
| 23 | return this.minHeap.peek(); |
| 24 | } |
| 25 | } |
| 1 | #include <queue> |
| 2 | #include <vector> |
| 3 | |
| 4 | using namespace std; |
| 5 | |
| 6 | class KthLargest { |
| 7 | private: |
| 8 | priority_queue<int, vector<int>, greater<int>> minHeap; |
| 9 | int k; |
| 10 | |
| 11 | public: |
| 12 | KthLargest(int k, vector<int>& nums) : k(k) { |
| 13 | for (int num : nums) { |
| 14 | minHeap.push(num); |
| 15 | } |
| 16 | while ((int)minHeap.size() > k) { |
| 17 | minHeap.pop(); |
| 18 | } |
| 19 | } |
| 20 | |
| 21 | int add(int val) { |
| 22 | minHeap.push(val); |
| 23 | if ((int)minHeap.size() > k) { |
| 24 | minHeap.pop(); |
| 25 | } |
| 26 | return minHeap.top(); |
| 27 | } |
| 28 | }; |
| 1 | class KthLargest { |
| 2 | constructor(k, nums) { |
| 3 | this.minHeap = []; |
| 4 | this.k = k; |
| 5 | for (const num of nums) { |
| 6 | this._heapPush(num); |
| 7 | } |
| 8 | while (this.minHeap.length > k) { |
| 9 | this._heapPop(); |
| 10 | } |
| 11 | } |
| 12 | |
| 13 | add(val) { |
| 14 | this._heapPush(val); |
| 15 | if (this.minHeap.length > this.k) { |
| 16 | this._heapPop(); |
| 17 | } |
| 18 | return this.minHeap[0]; |
| 19 | } |
| 20 | |
| 21 | _heapPush(val) { |
| 22 | this.minHeap.push(val); |
| 23 | this._siftUp(this.minHeap.length - 1); |
| 24 | } |
| 25 | |
| 26 | _heapPop() { |
| 27 | const last = this.minHeap.pop(); |
| 28 | if (this.minHeap.length === 0) return last; |
| 29 | const returnItem = this.minHeap[0]; |
| 30 | this.minHeap[0] = last; |
| 31 | this._siftDown(0); |
| 32 | return returnItem; |
| 33 | } |
| 34 | |
| 35 | _siftUp(idx) { |
| 36 | let parent = Math.floor((idx - 1) / 2); |
| 37 | while (idx > 0 && this.minHeap[idx] < this.minHeap[parent]) { |
| 38 | [this.minHeap[idx], this.minHeap[parent]] = [this.minHeap[parent], this.minHeap[idx]]; |
| 39 | idx = parent; |
| 40 | parent = Math.floor((idx - 1) / 2); |
| 41 | } |
| 42 | } |
| 43 | |
| 44 | _siftDown(idx) { |
| 45 | const n = this.minHeap.length; |
| 46 | while (true) { |
| 47 | let left = 2 * idx + 1; |
| 48 | let right = 2 * idx + 2; |
| 49 | let smallest = idx; |
| 50 | if (left < n && this.minHeap[left] < this.minHeap[smallest]) { |
| 51 | smallest = left; |
| 52 | } |
| 53 | if (right < n && this.minHeap[right] < this.minHeap[smallest]) { |
| 54 | smallest = right; |
| 55 | } |
| 56 | if (smallest === idx) break; |
| 57 | [this.minHeap[idx], this.minHeap[smallest]] = [this.minHeap[smallest], this.minHeap[idx]]; |
| 58 | idx = smallest; |
| 59 | } |
| 60 | } |
| 61 | } |
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
| 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] |
Objective
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.
Key Insight
Adding a new element to the heap takes O(log k) time. If the heap grows beyond size k, removing the smallest element ensures only the top k largest elements remain. Returning the root after each addition provides the current kth largest element efficiently.
Interview Quick-Check
Core Logic
Push the new value into the min-heap and pop the smallest element if the heap size exceeds k, ensuring the heap always contains the k largest elements.
State & Boundaries
Return the root of the heap after adjustment, which is the kth largest element.
Common Pitfalls & Bugs
Not popping when the heap size exceeds k leads to incorrect kth largest values and increased space usage.
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.
return self.min_heap[0]Return the root of the min-heap, representing the kth largest element.
The root of the min-heap is the smallest among the top k elements, which by definition is the kth largest element in the stream, providing O(1) retrieval.
heapq.heappop(self.min_heap)Pop the smallest element from the heap.
Removing the smallest element maintains the heap size constraint and preserves the heap property, ensuring correctness of kth largest tracking.
heapq.heappush(self.min_heap, val)Add the new value val to the min-heap.
Inserting val into the heap maintains the dynamic stream of elements, allowing the heap to update the top k largest elements efficiently.
heapq.heappop(self.min_heap)Pop the smallest element from the heap.
Removing the smallest element after exceeding size k discards elements that cannot be among the top k largest, preserving heap correctness.
self.min_heap = numsAssign the input list nums to the instance variable min_heap.
Using the input list directly as the heap storage avoids extra space allocation and allows in-place heap operations, improving efficiency.
self.k = kStore the integer k as an instance variable.
Preserving k allows the add method to maintain the heap size constraint consistently across calls.
Test Your Understanding
Why does maintaining a min-heap of size k guarantee that the root is the kth largest element in the stream?
Because the min-heap always contains the k largest elements seen so far, the smallest among them (the root) is the kth largest overall; any smaller elements are discarded, ensuring the heap accurately tracks the top k.
Related Problems
Heaps 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