LC 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.

Links

Question here and solution here

Solution

concept

The main trick here is to use a min heap instead of a max heap. This is because we only want the kth smallest element and if we keep a min heap with max size k then the top element will be the answer. This tricks only works because we do need to remove element, elements are just added in.

code

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.min_heap = []
        self.k = k
        for n in nums:
            heapq.heappush(self.min_heap, n)
        
        while len(self.min_heap) > k:
            heapq.heappop(self.min_heap)
        

    def add(self, val: int) -> int:
        heapq.heappush(self.min_heap, val)
        while len(self.min_heap) > self.k:
            heapq.heappop(self.min_heap)

        return self.min_heap[0]
        

# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)

from typing import List
import heapq

class KthLargest:
    """
    neetcode solution
    time: O(nlogn) for init, O(logn) for add
    """
    def __init__(self, k: int, nums: List[int]):
        # implement minHeap with k size
        # i.e. minHeap with k largest int
        self.minHeap = nums
        self.k = k
        # O(n)
        heapq.heapify(self.minHeap)
        # O(nlogn)
        while len(self.minHeap) > k:
            heapq.heappop(self.minHeap)

    def add(self, val: int) -> int:
        heapq.heappush(self.minHeap, val)
        # edge case if the initial list is less tha k
        # only pop if we have more than k element
        if len(self.minHeap) > self.k:
            heapq.heappop(self.minHeap)
        return self.minHeap[0]

Complexity

time: $O(\log n)$
space: $O(1)$