BinarySearch 0932 Kth Largest Numbers From Stream

Problem statement

https://binarysearch.com/problems/Kth-Largest-Numbers-From-Stream/

Solution

Equal to Leetcode 0703 Kth Largest Element in a Stream.

Complexity

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

Code

class KthLargestStream:
    def __init__(self, nums, k):
        self.heap = nums[:k+1]
        heapify(self.heap)
        self.k = k+1
        for num in nums[k+1:]:
            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]