Difficulty: Easy Pattern: Heap / Priority Queue LeetCode: https://leetcode.com/problems/kth-largest-element-in-a-stream/
PriorityQueue.Design a class that receives an integer k and an integer array nums (the initial stream), and
has a method add(val) that appends val to the stream then returns the k-th largest element
in the stream. If two elements are equal they each count separately (so [5,5,5] with k=2 has 2nd
largest = 5).
Signature:
class KthLargest:
constructor(int k, int[] nums)
int add(int val)
Examples:
Input: k = 3, nums = [4,5,8,2]
add(3) -> 4
add(5) -> 5
add(10) -> 5
add(9) -> 8
add(4) -> 8
Output: [4, 5, 5, 8, 8]
Input: k = 1, nums = []
add(3) -> 3
add(5) -> 5
Output: [3, 5] # k=1 means "the single largest so far"
Imagine you are a hiring manager building a shortlist of the K best applicants out of a tall
pile of resumes, and you read them one at a time. Once your shortlist has K people, every new
applicant only gets in if they are better than the worst person currently on the list. If they
are, you drop the worst and slot the new one in. The person you keep an eye on – the one who would
be dropped next – is the weakest of your current K. This problem is exactly that, where “best”
means “largest number” and new resumes arrive one add at a time.
To make “drop the worst of the current K” cheap we use a min-heap – a container whose top is always the smallest of the items inside it. We keep the heap at exactly size K and store the K largest values seen so far. Because the heap’s top is the smallest of those K, the top is the K-th largest overall – and when a new value arrives that beats it, we pop the top (the weakest survivor) and push the new value. This is the size-K heap trick: keep only K items; when full, evict the worst. Each add costs O(log k), no matter how big the stream has grown.
The reversal here is the only thing that takes a moment: for “K-th largest” we use a min-heap (min on top), not a max-heap. The reason is exactly the resume analogy – the value we want to throw away is the smallest of the K survivors, so that value has to be the one sitting on top ready to pop. A max-heap would put the largest survivor on top, and polling it would throw away the very value we want to keep.
Smallest trace. Take k = 3 and the stream 4, 5, 8, 2 as the initial values, then add(3):
{4}, size 1, not over 3 yet.{4, 5} (top is 4).{4, 5, 8} (top is 4). Size = 3 = K, full now.{2, 4, 5, 8} (size 4, one too many). Pop the top (2). Heap is back to
{4, 5, 8}, top = 4. So the 3rd largest of [4,5,8,2] is 4.add(3): push 3 -> {3, 4, 5, 8} (size 4), pop top (3) -> {4, 5, 8}, top = 4. Return 4.The constructor does steps 1-5 once for each starting value; add repeats steps 6 on every call.
This is the same size-K heap trick used by 0215 - Kth Largest Element in an
Array (the same question on a fixed array instead of a
stream) and 0347 - Top K Frequent (the heap stores value-frequency
pairs).
Pause and pick before expanding. A wrong first guess teaches more than a fast right one.
Q1 (recall). For “k-th largest” using a heap of size k, which heap puts the answer at the root?
Q2 (comprehend). Why is each add O(log k) and not O(log N), even after the stream has grown huge?
constructor(k, nums):
store k
create an empty min-heap
for each value x in nums:
push x into the heap
if heap size > k:
remove the smallest from the heap
function add(val):
push val into the heap
if heap size > k:
remove the smallest from the heap
return the smallest element now in the heap # that is the k-th largest
import java.util.PriorityQueue;
class KthLargest {
private final PriorityQueue<Integer> heap;
private final int k;
public KthLargest(int k, int[] nums) {
this.k = k;
this.heap = new PriorityQueue<>();
for (int x : nums) {
heap.offer(x);
if (heap.size() > k) {
heap.poll();
}
}
}
public int add(int val) {
heap.offer(val);
if (heap.size() > k) {
heap.poll();
}
return heap.peek();
}
}
The heap is a PriorityQueue<Integer>, which is a min-heap by default – exactly what we need,
because the root must be the smallest of the k largest (the element we would discard, and also the
answer). offer pushes and poll removes the root; both are O(log k) since the heap never exceeds
k+1 elements. The constructor and add share the identical “push then trim to k” step, which keeps
the invariant heap.size() <= k after every operation. peek at the end returns the k-th largest
in O(1). The guarantee “at least k elements exist when you search” means peek is never called on
an empty heap.
Time: O(N log k) to build, O(log k) per add -- each push/poll is bounded by heap height log k
Space: O(k) -- the heap holds at most k elements
Trace k = 3, nums = [4,5,8,2], then add(3) and add(10).
Constructor (build the initial heap of size 3):
| Step | x | heap after push | size > 3? | heap after poll | heap (sorted view) |
|---|---|---|---|---|---|
| 1 | 4 | {4} | no | {4} | [4] |
| 2 | 5 | {4,5} | no | {4,5} | [4,5] |
| 3 | 8 | {4,5,8} | no | {4,5,8} | [4,5,8] |
| 4 | 2 | {2,5,8,4} | yes | {5,8,4} | [4,5,8] |
After construction the root = 4, which is already the 3rd largest of [4,5,8,2].
add(3):
| Action | heap state | result |
|---|---|---|
| push 3 | {3,4,8,5} | - |
| size 4 > 3 poll | {4,5,8} | - |
| peek | root = 4 | 4 |
add(10):
| Action | heap state | result | |
|---|---|---|---|
| push 10 | {5,10,8,5} | - | (the two 5s are 4-then-5) |
| size 4 > 3 poll | {5,8,10} | - | |
| peek | root = 5 | 5 |
The survivors after add(10) are the three largest values 5,8,10; the root 5 is the 3rd largest.
Q1 (apply). Take k = 2, nums = [1,2,3], then call add(0). What does add(0) return?
Q2 (analyze). What goes wrong if add polls the root BEFORE pushing the new value?
Q3 (transfer). If the class had to return the k-th SMALLEST instead of the k-th largest, what is the one change that does it?
size > k check grows the heap to the whole
stream, turning O(N log k) into O(N log N) and breaking the size invariant.add. That would drop the wrong element on the very call that brings
the heap from k to k+1. Always push first, then trim.peek() without handling the constructor having fewer than k starting elements. The
problem guarantees k elements exist at query time, but the heap legitimately starts smaller; the
code is correct as long as you let add fill it up before peeking.[1,2,3] with k=2 answers 2 (2nd
largest), not 1.