Kth Largest Element in a Stream - LeetCode
Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element. Implement KthLargest class:
KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums. int add(int val) Appends the integer val to the stream and returns the element representing the kth largest element in the stream. Example 1: Input [“KthLargest”, “add”, “add”, “add”, “add”, “add”] [[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]] Output [null, 4, 5, 5, 8, 8] Explanation KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]); kthLargest.add(3); // return 4 kthLargest.add(5); // return 5 kthLargest.add(10); // return 5
- code
class KthLargest {
private PriorityQueue<Integer> pq;
private int qsize;
public KthLargest(int k, int[] nums) {
pq = new PriorityQueue<>();
qsize = k;
for (int i=0; i< nums.length; i++){
pq.offer(nums[i]);
if (pq.size() > k){
pq.poll();
}
}
}
public int add(int val) {
pq.offer(val);
if (pq.size() > qsize){
pq.poll();
}
return pq.peek();
}
}
- code
import heapq
class KthLargest:
"""
The idea is to ALWAYS maintain a MIN heap with only K elements
- in this case, the K-the largest element (in the stream)
- will always be at the root position
"""
def __init__(self, k: int, nums: List[int]):
self.heap = []
self.k = k
for num in nums:
self.add(num) # add elements using the function below
def add(self, val: int) -> int:
heapq.heappush(self.heap, val)
# if after adding the new item causes
# the heap size to increase beyond k,
# then pop out the smallest element
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0] # the root element