Kth Largest Element in a Stream
Understanding the Problem
We need a data structure that handles a growing collection of numbers and answers "what is the kth largest element right now?" every time a new number arrives. The stream only grows (we never remove elements), and we return the answer after each insertion.
We don't need to track all elements or maintain a fully sorted list. We only care about the top k elements at any given time. If we maintain only the k largest elements seen so far, the smallest among them is the kth largest overall, which is the answer.
Key Constraints:
0 <= nums.length <= 10^4→ The initial array can be empty, so the heap (or list) may start with fewer than k elements. It is guaranteed to reach at least k before any query needs an answer.- With up to
10^4initial elements plus10^4addcalls, we process around 20,000 elements total. An O(n) per-call solution gives roughly O(n^2) overall work, so the per-call cost should be sublinear. -10^4 <= val <= 10^4→ Values can be negative, so the comparison logic must not assume all elements are positive. The values fit comfortably in a 32-bit integer, so overflow is not a concern.
Approach 1: Sort Each Time
Intuition
Every time add is called, insert the new element into the list, sort the entire list in ascending order, and return the kth largest element. After sorting, the largest element is at index n - 1, the second largest at n - 2, so the kth largest is at index n - k.
This is straightforward to reason about: the list stays sorted, and the answer is always a fixed offset from the end.
Algorithm
- In the constructor, store the initial array and the value of
k. - When
add(val)is called, appendvalto the list. - Sort the entire list.
- Return the element at index
list.size() - k(the kth largest from the end).
Example Walkthrough
Code
- Time Complexity: O(n log n) per
addcall, where n is the current size of the list. Sorting takes O(n log n), and we do this on every call. - Space Complexity: O(n), where n is the total number of elements added. We store every element we've ever seen.
Sorting the entire list on every call does more work than the problem requires. We only need the kth largest, not a full ordering of every element. The next approach keeps only the k largest elements in a structure that exposes their minimum in O(1) and updates in O(log k).
Approach 2: Min-Heap of Size K
Intuition
We only care about the k largest elements at any point. An element smaller than all of the current top k can never be the kth largest, now or later (the stream only grows, so the top k can only get larger), so it can be discarded the moment it loses.
A min-heap of size k holds exactly those k largest elements. The heap's root is the smallest of them, which is the kth largest overall. When a new value arrives, push it onto the heap, then if the heap exceeds size k, pop the minimum. The element that gets popped is the one that no longer belongs in the top k.
The size-k min-heap maintains an invariant: after each operation it contains exactly the k largest elements seen so far. Pushing then popping the minimum preserves this. If the new value is among the top k, it stays and the old kth largest is evicted; if it is not, it becomes the new minimum and is evicted immediately, leaving the heap unchanged. Either way the root is the kth largest, so returning it answers the query.
Algorithm
- In the constructor, create a min-heap. Add all elements from
numsinto the heap. - While the heap size exceeds k, remove the minimum. This trims the heap to only the k largest elements.
- When
add(val)is called, pushvalonto the heap. - If the heap size exceeds k, pop the minimum.
- Return the heap's root (the kth largest element).
Example Walkthrough
Code
- Time Complexity: O(n log k) for the constructor (adding n initial elements to a heap of size k), and O(log k) per
addcall. Each heap push/pop on a heap of size k takes O(log k). - Space Complexity: O(k). We only store k elements in the heap at any time, regardless of how many total elements have been added.