LeetCode ยท #703

Kth Largest Element in a Stream Solution

Design a class that finds the kth largest element in a stream of numbers. Each call to add returns the current kth largest.

01
Section One ยท Problem

Problem Statement

๐Ÿท๏ธ

Difficulty

Easy

๐Ÿ”—

LeetCode

Problem #703

๐Ÿ—๏ธ

Pattern

Heap โ€” min-heap of size k

Design a class KthLargest that accepts an integer k and an initial array of integers. Implement an add(val) method that adds val to the stream and returns the element that is the kth largest in the collection so far.

Example
Input: k = 3, nums = [4, 5, 8, 2] add(3) โ†’ 4 // stream: [2,3,4,5,8] โ†’ 3rd largest = 4 add(5) โ†’ 5 // stream: [2,3,4,5,5,8] โ†’ 3rd largest = 5 add(10) โ†’ 5 // stream: [2,3,4,5,5,8,10] โ†’ 3rd largest = 5 add(9) โ†’ 8 // stream: [2,3,4,5,5,8,9,10] โ†’ 3rd largest = 8 add(4) โ†’ 8 // stream: [2,3,4,4,5,5,8,9,10] โ†’ 3rd largest = 8
Constraints
โ€ข 1 โ‰ค k โ‰ค 10โด โ€ข 0 โ‰ค nums.length โ‰ค 10โด โ€ข -10โด โ‰ค val โ‰ค 10โด โ€ข At most 10โด calls to add โ†‘ Stream grows unboundedly โ€” we shouldn't store everything
02
Section Two ยท Approach 1

Brute Force โ€” Sorted List O(n log n)

Keep all elements in a sorted list. On each add, insert the new value and re-sort (or insert at the correct position). The kth largest is at index length - k.

Why it works

Sorting guarantees the kth largest is always at a predictable index. Correct by definition โ€” after sorting, count k positions from the end.

Why we can do better
Problem: Re-sorting or finding the insertion point is O(n) per add. We only need the kth largest โ€” not the full sorted order. A min-heap of size k maintains exactly the top-k elements; its root is always the kth largest โ€” O(log k) per add.
Java โ€” Sorted List
import java.util.*; class KthLargest { List<Integer> stream = new ArrayList<>(); int k; public KthLargest(int k, int[] nums) { this.k = k; for (int n : nums) stream.add(n); Collections.sort(stream); } public int add(int val) { stream.add(val); Collections.sort(stream); // O(n log n) every time return stream.get(stream.size() - k); } }
Python โ€” Sorted List
class KthLargest: def __init__(self, k: int, nums: list[int]): self.k = k self.stream = sorted(nums) def add(self, val: int) -> int: self.stream.append(val) self.stream.sort() # O(n log n) every time return self.stream[-self.k]
MetricValue
Time (per add)O(n log n) โ€” re-sort on every insertion
SpaceO(n) โ€” stores all elements
03
Section Three ยท Approach 2

Min-Heap of Size k โ€” O(log k) per add

Maintain a min-heap of at most k elements โ€” the k largest values seen so far. The root (smallest in the heap) is always the kth largest overall. When a new value arrives: if the heap has fewer than k elements, add it; otherwise, only add it if it's larger than the root (then remove the root first).

๐Ÿ’ก Mental model: Imagine a VIP lounge with exactly k seats. A bouncer stands at the door. When someone new arrives, the bouncer checks if they outrank the weakest person inside (the root of your heap). If yes, the weakest person leaves and the newcomer sits down. The weakest person in the lounge is always the kth best in the entire crowd โ€” and you find them instantly because they're always near the door (the heap root).
Algorithm
  • Constructor: Add all initial elements to a min-heap. If heap size exceeds k, poll smallest until size = k.
  • add(val):
  • Offer val to the heap.
  • If heap size > k, poll the root (smallest).
  • Return heap.peek() โ€” the kth largest.
๐ŸŽฏ When to recognize this pattern:
  • Any problem asking for the kth largest/smallest in a dynamic or streaming context โ€” think min-heap of size k.
  • The signal is "maintain a running top-k." This pattern appears in LC 703, LC 215 (Kth Largest in Array), LC 973 (K Closest Points), and LC 347 (Top K Frequent).
  • For kth smallest, flip to a max-heap of size k.
Why min-heap, not max-heap?:
  • A min-heap of size k keeps the k largest elements.
  • The root is the smallest of those k โ€” which is exactly the kth largest overall.
  • If you used a max-heap, you'd need to store all n elements and can't efficiently discard the irrelevant small values.
04
Section Four ยท Trace

Visual Walkthrough

Trace with k = 3, nums = [4, 5, 8, 2], then add(3), add(5), add(10):

Min-Heap of size k=3 โ€” root is always the kth largest
CONSTRUCTOR: k=3, nums=[4, 5, 8, 2] Build heap with all 4 elements, then trim to size k=3 Add 4,5,8,2 โ†’ heap = [2,4,5,8] (size 4 > k=3) Poll root 2 โ†’ heap shrinks to size 3 Heap after init: 4 root (kth) 5 8 ADD OPERATIONS add(3): 3 < root(4) โ†’ ignore (doesn't belong in top-3) Heap unchanged: [4, 5, 8]. peek() = 4 โ†’ 4 add(5): 5 > root(4) โ†’ offer 5, heap size=4 โ†’ poll root 4 Heap: [5, 5, 8] โ€” old root 4 evicted, new root is 5 New 3rd largest is 5. peek() = 5 โ†’ 5 add(10): 10 > root(5) โ†’ offer 10, heap size=4 โ†’ poll root 5 Heap: [5, 8, 10] โ€” one of the 5s retained, new root still 5 3rd largest is 5. peek() = 5 โ†’ 5 add(9): 9 > root(5) โ†’ offer 9, heap size=4 โ†’ poll root 5 Heap: [8, 9, 10] โ€” 5 evicted, new root is 8 3rd largest is now 8. peek() = 8 โ†’ 8 Heap always holds the 3 largest values. Root = smallest of top-3 = kth largest. Returns: 4, 5, 5, 8, 8 โ€” all correct โœ“
Notice:
  • When add(3) is called, 3 is smaller than the root โ€” it can never be in the top-3, so we skip it.
  • Only values larger than the root get promoted into the heap, evicting the current root. The heap never grows beyond k=3.
05
Section Five ยท Implementation

Code โ€” Java & Python

Java โ€” Min-Heap of size k
import java.util.PriorityQueue; class KthLargest { PriorityQueue<Integer> heap; // min-heap int k; public KthLargest(int k, int[] nums) { this.k = k; heap = new PriorityQueue<>(); for (int n : nums) add(n); // reuse add logic } public int add(int val) { heap.offer(val); if (heap.size() > k) heap.poll(); // evict smallest โ€” keeps top-k return heap.peek(); // root = kth largest } }
Python โ€” Min-Heap of size k
import heapq class KthLargest: def __init__(self, k: int, nums: list[int]): self.k = k self.heap = nums[:] heapq.heapify(self.heap) while len(self.heap) > k: heapq.heappop(self.heap) # trim to size k def add(self, val: int) -> int: heapq.heappush(self.heap, val) if len(self.heap) > self.k: heapq.heappop(self.heap) # evict smallest return self.heap[0] # root = kth largest
06
Section Six ยท Analysis

Complexity Analysis

ApproachTime (per add)SpaceTrade-off
Sorted listO(n log n)O(n)Re-sort on every insertion
Binary search insertO(n)O(n)O(log n) to find position, O(n) to shift
Min-Heap size k โ† optimal O(log k) O(k) Only stores k elements. O(log k) per offer/poll.
Constructor cost:
  • The constructor is O(n log k) โ€” it adds n elements, each costing O(log k).
  • In Python, you can optimize to O(n) with heapify then trim, but the asymptotic bottleneck is the add calls over the stream's lifetime.
07
Section Seven ยท Edge Cases

Edge Cases & Pitfalls

CaseInputWhy It Matters
Empty initial arrayk=1, nums=[]First add must work. Heap starts empty, first value becomes root.
k equals array lengthk=3, nums=[1,2,3]All elements in heap. Root is the smallest = kth largest.
Duplicate valuesk=2, add(5), add(5)Heap can contain duplicates. Both 5s are valid top-k members.
Negative numbersk=1, nums=[-1,-2,-3]Min-heap works identically with negatives. Peek returns -1 (largest).
Value smaller than rootheap=[4,5,8], add(1)1 is offered, heap grows to 4, poll evicts 1 immediately. No change to top-k.
Monotonically increasing streamadd(1), add(2), add(3)...Heap root keeps rising. Every new value enters the heap, oldest root evicted.
โš  Common Mistake: Trying to optimize by checking if val > heap.peek() before offering โ€” but this fails when the heap has fewer than k elements (need to accept everything until full). The simpler pattern is: always offer, then poll if size exceeds k. This handles both under-full and full heaps uniformly.

โ† Back to Heap problems