Kth Largest Element in a Stream
Problem
Design a class that tracks the k-th largest element in a stream of integers. Each add(val) returns the k-th largest seen so far.
k = 3, nums = [4, 5, 8, 2]; add(3); add(5); add(10); add(9); add(4)4, 5, 5, 8, 8import heapq
class KthLargest:
def __init__(self, k, nums):
self.k = k
self.h = []
for x in nums: self.add(x)
def add(self, val):
if len(self.h) < self.k:
heapq.heappush(self.h, val)
elif val > self.h[0]:
heapq.heapreplace(self.h, val)
return self.h[0]// Sketch with sorted array stand-in for a min-heap of size k.
class KthLargest {
constructor(k, nums) { this.k = k; this.h = []; for (const x of nums) this.add(x); }
add(val) {
if (this.h.length < this.k) {
this.h.push(val); this.h.sort((a, b) => a - b);
} else if (val > this.h[0]) {
this.h[0] = val; this.h.sort((a, b) => a - b);
}
return this.h[0];
}
}class KthLargest {
int k;
PriorityQueue<Integer> h = new PriorityQueue<>();
public KthLargest(int k, int[] nums) { this.k = k; for (int x : nums) add(x); }
public int add(int val) {
if (h.size() < k) h.offer(val);
else if (val > h.peek()) { h.poll(); h.offer(val); }
return h.peek();
}
}class KthLargest {
int k;
priority_queue<int, vector<int>, greater<int>> h;
public:
KthLargest(int k, vector<int>& nums) : k(k) { for (int x : nums) add(x); }
int add(int val) {
if ((int)h.size() < k) h.push(val);
else if (val > h.top()) { h.pop(); h.push(val); }
return h.top();
}
};Explanation
We need the k-th largest value at any moment in a growing stream. The trick is to keep a min-heap holding exactly the k largest values seen so far — then the heap's smallest element (its root) is precisely the k-th largest.
Why a min-heap for a "largest" question? Because the smallest of the top k is the k-th largest. Keeping the heap at size k means the root is always the answer, and we never store more than we need.
On each add(val): if the heap has fewer than k items, just push val. Otherwise, compare val with the root. If val is bigger, it belongs in the top k, so we replace the root with it; if not, we ignore it. Either way we return the root.
Example: k=3 with starting nums [4,5,8,2] leaves the heap holding the three largest [4,5,8], root 4. Then add(3) → 3 isn't bigger than 4, answer stays 4. add(5) → 5 > 4, evict 4, heap is [5,5,8], answer 5. add(10) → 10 > 5, heap becomes [5,8,10], answer 5.
Each add touches the heap once, so it runs in O(log k) time, regardless of how long the stream gets.