Top-K / Heap — Kth Largest Element in a Stream
↗ LeetCode: Kth Largest Element in a Stream · Easy · pattern chapter →
Design a class KthLargest. Constructor takes k and an initial array; add(val) returns the k-th largest element after inserting val into the stream.
Example 1 —
KthLargest kthLargest = new KthLargest(3, [4,5,8,2]);
kthLargest.add(3); // returns 4
kthLargest.add(5); // returns 5
kthLargest.add(10); // returns 5
kthLargest.add(9); // returns 8
kthLargest.add(4); // returns 8Example 2 — KthLargest(1, []) then add(-3), add(-2), add(-4) → -3, -2, -2 (k=1 = running max). Example 3 — KthLargest(2, [1,1]) then add(1) → 1 (all equal).
Constraints — 1 ≤ k ≤ 10⁴; 0 ≤ nums.length ≤ 10⁴; -10⁴ ≤ vals ≤ 10⁴; at most 10⁴ add calls. Naive re-sort on each add is O(n log n) per call = O(n² log n) total = 10¹⁰ ops = TLE. Min-heap of size k is O(log k) per add = 10⁴ · 14 = 10⁵ ops.
Interview timer00:00
Approach 1 — Re-sort on every add
Intuition. Store all values; sort on each add; return arr[n - k].
Complexity — Time O(n log n) per add; too slow.
Approach 2 — Min-heap of size k
Insight from re-sort. We only ever care about the k largest values seen so far. A min-heap of size k whose root is the k-th largest.
- On
add(val): push. If heap size > k, poll the smallest. The root is now the k-th largest.
class KthLargest {
PriorityQueue<Integer> pq;
int k;
KthLargest(int k, int[] nums) {
this.k = k;
pq = new PriorityQueue<>();
for (int x : nums) add(x);
}
public int add(int val) {
pq.offer(val);
if (pq.size() > k) pq.poll();
return pq.peek();
}
}Complexity — Time O(log k) per add; Space O(k).
Approach 3 — Sorted TreeMap (with counts if duplicates allowed)
Insight from heap. A TreeMap<Integer, Integer> (value → count) also supports "kth largest" in O(log k). Not simpler than a heap, but useful if you need range queries too.
Complexity — Same as heap.
Try it yourself
Complexity summary
| Approach | Time per add | Space | Interview grade |
|---|---|---|---|
| Re-sort | O(n log n) | O(n) | baseline; TLE at 10⁴ adds |
| Min-heap of size k | O(log k) | O(k) | canonical |
| TreeMap | O(log k) | O(n) | overkill unless range queries needed |
When to use which
- Streaming k-th largest → min-heap of size k.
- Streaming k-th smallest → max-heap of size k.
- Both k-th largest AND range queries → TreeMap.
kvery large (close to n) → maintain running sort; heap advantage vanishes.
Related problems
- Top K Frequent Elements
- K Closest Points to Origin
- Kth Largest Element in an Array — offline sibling
- Find Median from Data Stream — two-heap for streaming median