Kth Largest in Stream
Kth Largest Element In An Array: class that, on each add(val), returns the k-th largest value seen so far.
- 0 <= nums.length <= 10⁴
- 1 <= k <= nums.length + 1
- -10⁴ <= nums[i] <= 10⁴
- -10⁴ <= val <= 10⁴
- At most 10⁴ calls will be made to add.
Intuition
This is the kth largest in stream problem — a class that is constructed once and then answers repeated add() calls, each returning the k-th largest value seen so far. The word stream is what makes it different from the array version: you do not have all the data up front, and you will be asked the same question many times.
Re-sorting on every call would cost O(n log n) each time. But look closely at what the question actually needs: only the top k values matter. The 500th largest value can never be the answer while k is 3, so there is no reason to keep it around.
So maintain exactly the k largest values seen so far, in a structure that makes the smallest of them cheap to find — because the smallest of the top k is the k-th largest. That is a min-heap of size k, and its root is the answer at all times:
- The heap's root is the answer, available in O(1) on every call.
Each add() then becomes: push the new value, and if the heap has grown to k+1 elements, pop the minimum to shrink it back. If the incoming value was too small to belong in the top k, it is the very thing that gets popped, and the heap is unchanged. Notice that a min-heap is used to track the largest values — that inversion is the part worth pausing on, and it is the same idea behind the size-k heap in Top K Frequent Elements.
Same min-heap-of-size-k idea as the static k-th largest, but now it has to survive repeated insertions. Capping the heap at k after every add means the root is permanently the answer, so each query is O(1) and each insert O(log k). Whenever a design question asks for a running order statistic, this is the structure.
Approach
Before reading on: price up what the direct approach costs here, then ask why only the extreme element matters and the rest need no ordering. Aim for O(log k) per add time and O(k) space.
Hold the invariant, not the data
The object stores only the k largest values seen so far, never the full stream. Everything else is discarded permanently. This is what keeps memory at O(k) even when millions of values arrive.
Use a min-heap so the answer sits on top
A min-heap keeps its smallest element at the root. Among the k largest values, the smallest is exactly the k-th largest overall — so the root is the answer and reading it costs O(1). Using a max-heap here would put the wrong end within reach.
Seed the heap in the constructor
Push the initial array, then pop until the size is k. Alternatively heapify the whole array in O(n) and pop down — worth doing when the seed array is much larger than k, since heapify beats n individual pushes.
On add, push then trim
Push the new value unconditionally, then if the size exceeds k, pop the minimum. A value too small to belong in the top k gets pushed and immediately popped, leaving the heap exactly as it was — no branching or special-casing needed.
Return the root
After the trim, the heap holds precisely the k largest values seen so far, and its root is the smallest of them. Return it. The invariant is restored before every return, so consecutive calls stay correct without any extra bookkeeping.
Why a sorted list is worse
Keeping a sorted array would give O(1) reads, but each insertion costs O(n) for the shifting. The heap trades a slightly slower read for O(log k) insertion, which is the operation that runs on every element of the stream and therefore dominates the total cost.
Cost per operation
Construction is O(n log k) with repeated pushes, or O(n) with heapify plus O((n − k) log k) of popping. Each add() is O(log k) and each read is O(1). Space is O(k) regardless of how long the stream runs, which is the property that makes this viable on unbounded input.
Solution & live demo
Common pitfalls
Trimming the heap only in the constructor
def add(self, val):
heapq.heappush(self.heap, val)
return self.heap[0]heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]The heap grows past k and its root drifts down to the overall minimum, so the answer becomes the smallest value ever added rather than the k-th largest. The cap must be re-applied on every insertion.
Rejecting small values instead of pushing then popping
if val > self.heap[0]:
heapq.heapreplace(self.heap, val)
return self.heap[0]heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)During warm-up the heap may hold fewer than k elements, and then even a small value belongs in it — the guard wrongly discards it and heap[0] crashes on an empty heap. Push-then-trim is correct in both phases.
Returning the largest element
return max(self.heap)
return self.heap[0]
The heap holds exactly the k largest values, so its smallest member is the k-th largest overall. Taking the maximum returns the single biggest element instead.
Edge cases
Heap fills up over the first adds; problem guarantees k values exist before queries matter.
Kept — k-th largest counts repeats.