Problem
We need to design a class that can continuously return the kth largest element from a stream of numbers.
The stream starts with an initial list nums, and new values are added one by one using the add() method.
Duplicates count as separate values.
For example, the 2nd largest element in:
[1, 2, 3, 3]is 3, because after sorting:
[1, 2, 3, 3]the two largest values are 3 and 3.
Brute Force Approach: Sort Every Time
A simple way to solve the problem is to store all numbers in a list. Every time we add a new number, we sort the list and return the kth largest element.
from typing import List
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.nums = nums
self.nums.sort()
def add(self, val: int) -> int:
self.nums.append(val)
self.nums.sort()
return self.nums[-self.k]How It Works
If we want the kth largest number, we can sort the array in ascending order and access:
self.nums[-self.k]For example, if:
nums = [1, 2, 3, 3, 5]
k = 3then:
nums[-3] == 3So the 3rd largest element is 3.
Complexity
Let n be the number of elements currently in the stream.
The constructor sorts the initial list:
Time: O(n log n)Each add() call appends a value and sorts again:
Time: O(n log n)This works, but it is inefficient because sorting the entire list every time is unnecessary.
Optimized Approach: Min Heap
A better solution is to keep only the largest k elements seen so far.
To do this, we use a min heap.
The smallest element inside this heap will always be the kth largest element overall.
Why?
Because the heap stores the top k largest elements. Among those k largest elements, the smallest one is exactly the kthlargest.
Heap Solution
import heapq
from typing import List
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.minHeap = nums
self.k = k
heapq.heapify(self.minHeap)
while len(self.minHeap) > k:
heapq.heappop(self.minHeap)
def add(self, val: int) -> int:
heapq.heappush(self.minHeap, val)
if len(self.minHeap) > self.k:
heapq.heappop(self.minHeap)
return self.minHeap[0]How the Heap Approach Works
We maintain a min heap of size at most k.
When a new value is added:
- Push the new value into the heap.
- If the heap size becomes greater than
k, remove the smallest value. - The root of the heap is the
kthlargest value.
The key idea is that we do not need to store every number in sorted order. We only care about the largest k numbers.
Example
Suppose:
k = 3
nums = [1, 2, 3, 3]After heapifying and keeping only 3 elements, the heap contains the 3 largest values:
[2, 3, 3]The smallest value in this heap is 2, so the 3rd largest is 2.
Now we call:
add(3)The heap becomes:
[2, 3, 3, 3]Since the heap size is greater than k, remove the smallest value:
[3, 3, 3]Now the 3rd largest is:
3So we return:
3Complexity
Let n be the length of the initial list.
The constructor heapifies the list:
O(n)Then it removes extra elements until only k remain:
O((n - k) log n)Each add() operation pushes one value and may pop one value:
O(log k)The heap never grows larger than k + 1, so the space complexity is:
O(k)Final Thoughts
The brute force solution is easier to understand because it simply sorts the list every time. However, it does more work than needed.
The heap solution is better because it focuses only on the largest k values. Since the smallest value among those k values is the answer, we can return the top of the min heap in constant time.
This makes the min heap approach the preferred solution for this problem.
