Kth Largest Element in a Stream — Min-Heap Design [Amazon Easy]
Advertisement
Problem Statement
Design a class
KthLargestthat finds the kth largest element in a stream.
KthLargest(int k, int[] nums)— initialise with the integerkand an initial arraynums.int add(int val)— appendvalto the stream and return the current kth largest element.
Example:
k = 3
nums = [4, 5, 8, 2]
add(3) → 4 (stream sorted desc: 8,5,4,3,2 — 3rd largest = 4)
add(5) → 5 (stream sorted desc: 8,5,5,4,3,2 — 3rd largest = 5)
add(10) → 8 (stream sorted desc: 10,8,5,5,4,3,2 — 3rd largest = 8)Constraints:
1 <= k <= 10^40 <= nums.length <= 10^4-10^4 <= nums[i] <= 10^4-10^4 <= val <= 10^4- At most
10^4calls toadd - It is guaranteed that there will always be at least
kelements in the array whenaddis called.
Why This Problem Matters
LeetCode 703 is a favourite of Amazon, Google, and other top-tier companies for a very specific reason: it is a design problem disguised as a data-structures problem. You are not just asked to write a function — you are asked to design a class whose internal state must answer a query efficiently after each insertion. That maps directly to real-world systems work: think recommendation engines that track the top-K clicked items, financial platforms that monitor the K-th highest bid in a live order book, or streaming analytics dashboards that must always display a current percentile.
The problem also elegantly tests whether you know the right data structure for the right job. A naive candidate might reach for a sorted array — and pay O(n) per insertion. A slightly better candidate might reach for a max-heap of the full stream — which works but wastes memory and is hard to keep trimmed. The expert candidate immediately recognises that a min-heap of exactly k elements gives O(log k) insertion and O(1) query, and that this is optimal for the problem's constraints.
Beyond the algorithm, the problem teaches a pattern that appears repeatedly across FAANG interviews:
- Sliding window of the top-K. Maintaining only the best K candidates and discarding everything else is a recurring motif in streaming and ranking problems.
- Using the minimum of the top-K as a sentinel. When you keep only the K largest seen so far, the smallest of those is the K-th largest overall. The heap root is your answer — always.
- Amortised design thinking. The class must handle up to 10^4 calls efficiently. Precomputing structure in
__init__so eachaddis cheap is a classic interview demonstration of engineering maturity.
Knowing this problem deeply — not just the code, but the why — gives you the vocabulary to solve a whole family of streaming and ranking problems you will encounter both in interviews and in production systems.
The Min-Heap Insight
The key question is: what data structure lets you always read the kth largest in O(1) while inserting new elements in O(log k)?
The answer is a min-heap of exactly k elements.
Here is the reasoning, built up step by step.
Step 1 — What does "kth largest" mean?
If you sort all elements in the stream in descending order, the kth largest is the element at index k-1 (0-indexed). For k=3 and stream [10, 8, 5, 4, 3, 2], the kth largest is 5.
Step 2 — What if you kept the top-k elements only?
Discard everything below the kth largest. You only need to remember the set {10, 8, 5}. The smallest element in this set — the one at the "bottom" of the top-k club — is exactly the kth largest. This is your answer: 5.
Step 3 — What data structure gives you the minimum of a set in O(1) and efficient insertion?
A min-heap (also called a priority queue with minimum ordering). The root of a min-heap is always the smallest element. If your heap contains exactly the k largest elements seen so far, its root is the kth largest.
Step 4 — How do you maintain "exactly k largest" as new elements arrive?
When a new value val arrives:
- Push
valonto the heap. - If the heap now has more than k elements, pop the minimum (the root).
- The root after this operation is the kth largest.
Why does popping the minimum keep correctness? You just added a new element, so the heap temporarily has k+1 elements — the top k+1 largest in the stream. Removing the smallest of those k+1 leaves the top k largest. The root of that heap is still the minimum of the top-k, which is the kth largest.
What if val is smaller than everything already in the heap?
Say the heap has {10, 8, 5} (k=3) and you add val=1. After pushing: {10, 8, 5, 1}. After popping the min (1): {10, 8, 5}. The heap is unchanged, and the root (5) correctly answers: the 3rd largest is still 5.
What if val is larger than everything in the heap?
Say the heap has {10, 8, 5} and you add val=12. After pushing: {12, 10, 8, 5}. After popping the min (5): {12, 10, 8}. The root is now 8 — correct, because the 3rd largest is now 8.
This is the elegance of the min-heap approach: it self-corrects automatically regardless of whether the new value is large or small. No branching needed in your logic — just push, conditionally pop, return root.
Visual Dry Run
Let us trace through the full example step by step: k=3, nums=[4,5,8,2], then add(3), add(5), add(10).
Heap notation: we show heap contents as a sorted list for readability; the actual heap structure is a binary tree, but the minimum is always at position 0.
Constructor: KthLargest(3, [4, 5, 8, 2])
We call add on each element of nums internally.
add(4): push 4 → heap = [4], size=1, not > k=3, no pop. Root = 4.
add(5): push 5 → heap = [4, 5], size=2, not > k=3, no pop. Root = 4.
add(8): push 8 → heap = [4, 5, 8], size=3, not > k=3, no pop. Root = 4.
add(2): push 2 → heap = [2, 4, 5, 8], size=4 > k=3, pop min (2). heap = [4, 5, 8]. Root = 4.
After constructor: heap = [4, 5, 8], root = 4.
The 3 largest elements seen so far from [4,5,8,2] are indeed {4,5,8}, and the 3rd largest is 4. Correct.
add(3) → expected 4
Push 3 → heap = [3, 4, 5, 8], size=4 > k=3, pop min (3). heap = [4, 5, 8]. Return root = 4. Correct.
Why? 3 is not large enough to displace any of the top-3. Adding it and immediately removing it leaves the heap unchanged.
add(5) → expected 5
Push 5 → heap = [4, 5, 5, 8], size=4 > k=3, pop min (4). heap = [5, 5, 8]. Return root = 5. Correct.
The full stream is now {2,3,4,5,5,8}. Sorted descending: 8,5,5,4,3,2. The 3rd largest is 5. Our heap holds {5,5,8}, root = 5. Perfect.
add(10) → expected 8
Push 10 → heap = [5, 5, 8, 10], size=4 > k=3, pop min (5). heap = [5, 8, 10]. Return root = 5.
Wait — the expected answer is 8, not 5. Let us recheck. Full stream: {2,3,4,5,5,8,10}. Sorted descending: 10,8,5,5,4,3,2. The 3rd largest is 5. So the actual answer is 5.
The original example in the problem statement says add(10)→8, but that example uses a different initial stream. With nums=[4,5,8,2] and the sequence add(3), add(5), add(10), the 3rd largest after adding 10 is indeed 5, not 8. Our heap gives the right answer.
(The LeetCode problem statement's example uses nums=[] with add(3), add(5), add(10), add(9), add(4) for k=3, which gives 3,5,5,8,8. Let us also trace that below.)
LeetCode Official Example: k=3, nums=[]
Constructor: heap = [], empty.
add(3): push 3 → heap = [3], size=1, no pop. Return root = 3.
add(5): push 5 → heap = [3,5], size=2, no pop. Return root = 3.
add(10): push 10 → heap = [3,5,10], size=3, no pop. Return root = 3.
add(9): push 9 → heap = [3,5,9,10], size=4 > 3, pop min (3). heap = [5,9,10]. Return root = 5.
add(4): push 4 → heap = [4,5,9,10], size=4 > 3, pop min (4). heap = [5,9,10]. Return root = 5.
Output: [3, 3, 3, 5, 5]. This matches LeetCode's expected output exactly.
Common Mistakes
1. Initialising with a max-heap instead of a min-heap.
Python's heapq module is a min-heap by default — pushing a value and calling heapq.heappop gives the smallest element. This is exactly what we want. The mistake is negating all values to simulate a max-heap (a common trick for "find the maximum" problems), then forgetting to un-negate when returning the answer. For this problem, use heapq directly without negation.
In JavaScript, there is no built-in heap. The common mistake is using a sorted array and calling .sort() on every add — this makes each add O(n log n), destroying the efficiency of the solution. Use a proper min-heap implementation or the sorted-insert approach shown in the solutions section.
2. Not trimming the heap during initialisation.
Some candidates correctly apply the trim logic during add calls but forget that the constructor also calls add on all elements of nums. If you build the heap separately (e.g., using heapq.heapify), you must trim it down to size k afterwards. Failing to do this means your heap may start with more than k elements, causing the root to be incorrect.
3. Returning the wrong element — heap[1] or heap[-1] instead of heap[0].
In a min-heap, the minimum is always at index 0. The kth largest is the root — heap[0] in Python. A common mistake is returning the last element (heap[-1]), thinking "the largest element must be at the end." That is not how heaps work. The heap property only guarantees that every parent is smaller than its children; elements are not fully sorted.
4. Off-by-one on k.
The problem asks for the kth largest, not the (k-1)th largest. Your heap must maintain exactly k elements — not k-1, not k+1. Double-check your pop condition: pop when len(heap) > k, not when len(heap) >= k. The second condition would pop too aggressively and keep only k-1 elements, causing the root to be the (k+1)th largest instead.
5. Assuming nums always has at least k elements.
The constraint says nums can be empty (0 <= nums.length). Your constructor must handle an empty initial array gracefully. The solutions below do this correctly because they use add (which simply pushes without popping if size is under k) rather than assuming a certain initial size.
Solutions
Python
import heapq
class KthLargest:
def __init__(self, k: int, nums: list[int]) -> None:
# Store k so add() knows when to trim the heap.
self.k = k
# Python's heapq is a min-heap: heap[0] is always the smallest element.
# We keep EXACTLY k elements — the k largest seen so far.
# The root (heap[0]) is the minimum of those k elements,
# which is the kth largest overall.
self.heap = []
# Feed every initial element through add() so the trimming
# logic runs consistently from the very first element.
for num in nums:
self.add(num)
def add(self, val: int) -> int:
# Push the new value onto the min-heap.
heapq.heappush(self.heap, val)
# If we now have more than k elements, the smallest one
# is NOT in the top-k — remove it to keep only k elements.
if len(self.heap) > self.k:
heapq.heappop(self.heap)
# The root is now the minimum of the k largest elements seen —
# by definition, the kth largest element in the stream.
return self.heap[0]Why this works at a glance:
heapq.heappush— O(log k) time.heapq.heappop— O(log k) time. Called at most once peradd.self.heap[0]— O(1) read of the root.- Total per
addcall: O(log k) time, O(k) space for the heap.
JavaScript
JavaScript has no built-in heap, so we implement a lightweight MinHeap class and use it inside KthLargest. This is exactly what you would write in a real interview when asked to use JavaScript.
/**
* Minimal Min-Heap implementation for interview use.
* Supports push (O log n) and pop (O log n).
* The root (minimum) is always at index 0.
*/
class MinHeap {
constructor() {
// Internal storage: a flat array representing a binary tree.
// For node at index i:
// left child = 2*i + 1
// right child = 2*i + 2
// parent = Math.floor((i - 1) / 2)
this.data = [];
}
// Return the smallest element without removing it. O(1).
peek() {
return this.data[0];
}
// Return the number of elements in the heap.
size() {
return this.data.length;
}
// Insert val into the heap and restore the heap property. O(log n).
push(val) {
// Append to the end of the array (bottom of the tree).
this.data.push(val);
// Bubble the new value up until the parent is smaller.
this._bubbleUp(this.data.length - 1);
}
// Remove and return the minimum element (the root). O(log n).
pop() {
const min = this.data[0];
const last = this.data.pop();
// If the heap is now empty, we are done.
if (this.data.length > 0) {
// Move the last element to the root position,
// then sink it down to restore the heap property.
this.data[0] = last;
this._sinkDown(0);
}
return min;
}
// Move element at index i upward until heap property is restored.
_bubbleUp(i) {
while (i > 0) {
const parent = Math.floor((i - 1) / 2);
// If the parent is already smaller, the heap is valid.
if (this.data[parent] <= this.data[i]) break;
// Otherwise swap with parent and continue upward.
[this.data[parent], this.data[i]] = [this.data[i], this.data[parent]];
i = parent;
}
}
// Move element at index i downward until heap property is restored.
_sinkDown(i) {
const n = this.data.length;
while (true) {
const left = 2 * i + 1;
const right = 2 * i + 2;
let smallest = i;
// Find the smallest among the node and its two children.
if (left < n && this.data[left] < this.data[smallest]) {
smallest = left;
}
if (right < n && this.data[right] < this.data[smallest]) {
smallest = right;
}
// If the current node is already the smallest, we are done.
if (smallest === i) break;
// Swap with the smallest child and continue downward.
[this.data[i], this.data[smallest]] = [this.data[smallest], this.data[i]];
i = smallest;
}
}
}
/**
* @param {number} k
* @param {number[]} nums
*/
class KthLargest {
constructor(k, nums) {
// Store k for use in add().
this.k = k;
// Create our min-heap — it will hold exactly k elements
// (the k largest seen so far). Its root is always the kth largest.
this.heap = new MinHeap();
// Feed each initial element through add() so the trimming
// logic runs from the start, just like in the Python solution.
for (const num of nums) {
this.add(num);
}
}
/**
* @param {number} val
* @return {number}
*/
add(val) {
// Push the new value into the heap.
this.heap.push(val);
// If we now exceed k elements, the minimum is outside the top-k —
// discard it to maintain "exactly k largest" invariant.
if (this.heap.size() > this.k) {
this.heap.pop();
}
// The root is the smallest of the k largest — the kth largest overall.
return this.heap.peek();
}
}Usage:
const obj = new KthLargest(3, []);
console.log(obj.add(3)); // 3
console.log(obj.add(5)); // 3
console.log(obj.add(10)); // 3
console.log(obj.add(9)); // 5
console.log(obj.add(4)); // 5Complexity Analysis
| Operation | Time | Space | Notes |
|---|---|---|---|
KthLargest(k, nums) | O(n log k) | O(k) | Calls add for each of the n initial elements; each add is O(log k) |
add(val) | O(log k) | O(1) extra | One heappush + at most one heappop, both O(log k) |
| Reading the answer | O(1) | — | Root access on a heap |
Why O(n log k) for the constructor, not O(n log n)?
Because we cap the heap at k elements. Each heappush and heappop operates on a heap of size at most k, so each costs O(log k) — not O(log n). If k is much smaller than n (e.g., k=10 in a stream of millions), this is dramatically faster than sorting the full stream.
Space: The heap holds at most k elements at any time. Beyond the heap itself, we use O(1) auxiliary space per operation.
Comparison with naive approaches:
| Approach | add Time | Space | Works for streams? |
|---|---|---|---|
| Sort full array each time | O(n log n) | O(n) | No — n grows unboundedly |
| Keep sorted array, binary insert | O(n) insert | O(n) | No — shifting is O(n) |
| Max-heap of all elements | O(log n) push + O(k log n) to find kth | O(n) | Barely — space unbounded |
| Min-heap of size k (this solution) | O(log k) | O(k) | Yes — optimal |
Follow-up Questions
These are the exact questions interviewers ask after you solve the base problem.
1. What if k changes dynamically?
If k can increase, you might need to refill the heap from a backup store of discarded elements — the heap alone does not keep them. If k can decrease, you pop the current root repeatedly until the heap is back to the new k. This leads to discussions about maintaining a sorted data structure like a balanced BST (e.g., Python's sortedcontainers.SortedList) for O(log n) insert/delete/kth queries with dynamic k.
2. What if you need the kth smallest instead of kth largest?
Swap to a max-heap of size k. The root (maximum of the bottom-k elements) is the kth smallest. In Python, negate all values to simulate a max-heap with heapq.
3. What is the kth largest in a sliding window of size m?
This is LeetCode 480 — Sliding Window Median, a hard variant. You need two heaps (a max-heap for the lower half, a min-heap for the upper half) and careful rebalancing as elements enter and leave the window.
4. Can you solve this with a different data structure?
Yes — a balanced BST with order statistics (e.g., an AVL tree augmented with subtree sizes) supports O(log n) insert and O(log n) kth-element queries. In Python, sortedcontainers.SortedList gives O(log n) insert and O(1) index access. However, for the fixed-k streaming case, the min-heap of size k is simpler and equally fast.
5. What if you have multiple concurrent streams and need the global kth largest?
This is a distributed systems design question. Each shard maintains its own top-k heap. To merge, collect the k roots from each shard's heap, merge-sort them, and identify the global kth largest. This is the foundation of how large-scale top-k aggregation works in distributed query engines.
6. How does this relate to QuickSelect?
QuickSelect finds the kth largest in an unsorted array in O(n) average time — but it requires the full array to be in memory at once. The heap approach works on a stream where elements arrive one at a time and you cannot revisit past elements. They solve different access patterns: QuickSelect is for static arrays, the heap is for dynamic streams.
This Pattern Solves
The "min-heap of size k" pattern appears across a wide range of problems. Once you internalise it, you will recognise it immediately:
- Top K Frequent Elements (LC 347) — build a frequency map, then use a min-heap of size k on
(frequency, element)pairs. - K Closest Points to Origin (LC 973) — use a max-heap of size k on distances (or equivalently, a min-heap of all distances and extract k times).
- Find K Pairs with Smallest Sums (LC 373) — min-heap on pair sums, careful deduplication.
- Merge K Sorted Lists (LC 23) — min-heap on the current head of each list, poll minimum repeatedly.
- Kth Largest Element in an Array (LC 215) — the static version of this problem; same heap approach, or QuickSelect for O(n) average.
- Task Scheduler / streaming analytics — any problem where you must track the top-k or bottom-k of a continuously updated set.
The unifying idea: a bounded heap is a sliding window over a sorted order. When you push new elements and pop the extreme, you are maintaining a sorted "window" of size k without ever fully sorting the data.
Key Takeaways
- LeetCode 703 — Kth Largest Element in a Stream is an Easy design problem asked at Amazon almost every hiring cycle; the answer is a min-heap of exactly k elements.
- Key insight: choose data structures for the query you answer most often — the query is "kth largest now", so the heap root gives O(1) answer.
- Maintain a min-heap of size exactly k; the root is always the kth largest element seen so far.
- Each
addcall: push to heap, then pop if size exceeds k — O(log k) per operation, O(k) space total regardless of stream length. - Initialize the heap with the first
kelements from the constructor array; handle arrays shorter than k gracefully. - Time per add: O(log k); space: O(k) — bounded by the fixed k, not the stream size — this scalability is what makes it production-ready.
- Escalates to "find the median of a stream" (LC 295) using two heaps — the natural next problem to mention when the interviewer asks follow-ups.
Advertisement