educative.io

More efficient solution

We don’t need to add everything to the heap blindly as we’re wasting computations like that.

import heapq

class KthLargestNumberInStream:
  def __init__(self, nums, k):
    # TODO: Write your code here
    self.k = k
    self.largest = []
    self.__startup(nums)

  def __startup(self, nums):
    count = 0
    for num in nums:
      # We havent added the minimum number of number inside the heap
      # then add them without checking anything
      if count < self.k:
        heapq.heappush(self.largest, num)
      elif self.largest[0] < num: 
        # The smallest number in the heap is smaller than the current number
        # so remove the smallest and add the current number
        heapq.heappop(self.largest)
        heapq.heappush(self.largest, num)

      count += 1


  def add(self, num):
    if self.largest[0] < num:
      heapq.heappop(self.largest)
      heapq.heappush(self.largest, num)

    return self.largest[0]

Instantiation of the class:
Time complexity: O(n log(k))
Space complexity : O(k)

Adding numbers:
Time complexity:
best case O(1) as we’ll just return smallest number in the heap if larger than the given value
worst case O(k) as we’'ll need to replace the number in the heap with the new one

Space complexity: O(1) no additional space used other than the original fixed k size heap

Correction adding a number in the worst case will take O(logk)