Problem statement

https://leetcode.com/problems/kth-largest-element-in-a-stream/

Solution

Similar to Problems 0973 K Closest Points to Origin and 0215 Kth Largest Element in an Array with exactly the same idea: here we heap min-heap with k biggest elements, then k-th largest element will be in the root of our heap. When we have new element we check if it is more than root of our heap, then we put this element to heap and remove smallest element from heap, so we will always keep it k elements.

Complexity

It is O(k + (n-k)*log k) for init and O(log k) for add. Space complexity is O(k).

Code

class KthLargest:
    def __init__(self, k, nums):
        self.heap = nums[:k]
        heapify(self.heap)
        self.k = k
        for num in nums[k:]:
            if num > self.heap[0]:
                heappushpop(self.heap, num)
                
    def add(self, val) :
        if len(self.heap) < self.k:
            heappush(self.heap, val)
        elif val > self.heap[0]:
            heappushpop(self.heap, val)
            
        return self.heap[0]