Kth Largest Element in a Stream
Advertisement
Problem
Design a class that adds numbers from a stream and returns the Kth largest element after each addition.
Key insight: Maintain a min-heap of size K. The minimum of the heap is the Kth largest.
Approach
- Init: heapify first K elements, push remaining one-by-one evicting if > K
- Add: push to heap, pop if size > K, return heap[0]
Solutions
// C — conceptual (use array-based min-heap)
// See C++ for clean implementation
// C++
class KthLargest {
priority_queue<int, vector<int>, greater<int>> minH;
int k;
public:
KthLargest(int k, vector<int>& nums) : k(k) {
for (int n : nums) add(n);
}
int add(int val) {
minH.push(val);
if ((int)minH.size() > k) minH.pop();
return minH.top();
}
};
// Java
class KthLargest {
PriorityQueue<Integer> minH = new PriorityQueue<>();
int k;
public KthLargest(int k, int[] nums) {
this.k = k;
for (int n : nums) add(n);
}
public int add(int val) {
minH.offer(val);
if (minH.size() > k) minH.poll();
return minH.peek();
}
}
// JavaScript (using sorted array for simplicity)
class KthLargest {
constructor(k, nums) {
this.k = k;
this.heap = [];
for (const n of nums) this.add(n);
}
add(val) {
this.heap.push(val);
this.heap.sort((a, b) => a - b);
if (this.heap.length > this.k) this.heap.shift();
return this.heap[0];
}
}
# Python
import heapq
class KthLargest:
def __init__(self, k, nums):
self.k = k
self.heap = []
for n in nums:
self.add(n)
def add(self, val):
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]
Complexity
- add(): O(log k)
- Space: O(k)
Key Insight
A min-heap of size K always has the Kth largest at the top (the minimum). Elements smaller than the Kth largest are evicted.
Advertisement