703. Kth Largest Element in a Stream

Question

You are part of a university admissions office and need to keep track of the kth highest test score from applicants in real-time. This helps to determine cut-off marks for interviews and admissions dynamically as new applicants submit their scores.
You are tasked to implement a class which, for a given integer k, maintains a stream of test scores and continuously returns the kth highest test score after a new score has been submitted. More specifically, we are looking for the kth highest score in the sorted list of all scores.
Implement the KthLargest class:

KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of test scores nums.
int add(int val) Adds a new test score val to the stream and returns the element representing the kth largest element in the pool of test scores so far.

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
kthLargest.add(9); // return 8
kthLargest.add(4); // return 8

Example 2:

Input:
["KthLargest", "add", "add", "add", "add"]
[[4, [7, 7, 7, 7, 8, 3]], [2], [10], [9], [9]]
Output: [null, 7, 7, 7, 8]
Explanation:
KthLargest kthLargest = new KthLargest(4, [7, 7, 7, 7, 8, 3]);
kthLargest.add(2); // return 7
kthLargest.add(10); // return 7
kthLargest.add(9); // return 7
kthLargest.add(9); // return 8
 
Constraints:

0 <= nums.length <= 104
1 <= k <= nums.length + 1
-104 <= nums[i] <= 104
-104 <= val <= 104
At most 104 calls will be made to add.

Prakash Sellathurai · Accepted Answer

# k -> [2.3.4] k =2 # 1 2 3 4 5 # add (6) -> 4 [1 2 3 4 5 6] # add(1) -> 5 [1 1 2 3 4 5 6] # > heap of size k min heap top of heap is kths largest element every add operation is O(1) # edge case if n < k return top # import heapq class KthLargest: def __init__(self, k: int, nums: List[int]): # 3, [4, 5, 8, 2] self.heap = [] self.k = k heapq.heapify(self.heap) for num in nums: # 4 , 5, 8, 2 heapq.heappush(self.heap, num) # [2, 4, 5, 8] if len(self.heap) > k: heapq.heappop(self.heap) # [4, 5, 8] # invariant make sure def add(self, val: int) -> int: #[4, 5, 8] val = 10 # val = heap[0] # val > heap[0] # heappush and pop # return val if len(self.heap) < self.k: heapq.heappush(self.heap, val) elif self.heap and val >= self.heap[0]: heapq.heappush(self.heap, val) # [4,5,8, 10] heapq.heappop(self.heap) # [5,8, 10] return self.heap[0] # Problem: # Kth Largest Stream # Invariant: # Heap contains k largest elements. # Failed assumption: # Always pop before push. # Final insight: # Only remove an element if the new value belongs in top k. # Your KthLargest object will be instantiated and called as such: # obj = KthLargest(k, nums) # param_1 = obj.add(val)