< Back





703. Kth Largest Element in a Stream

We're asked to design a class that will maintain the kth largest element seen in a stream. It also
returns the kth largest elements when new elements are added. To solve this, we just heapify the
input and, while the heap is greater than k, we continue to pop the heap. The kth largest element
is then the top of the heap.

Why? By default Python uses a min heap. If k == 3, and we keep the heap to size k, then the minimum
element in the heap is the kth largest element. We've already popped off all of the smaller elements.

The solution is as follows:


  class KthLargest:
      def __init__(self, k: int, nums: List[int]) -> None:
          self.k = k
          self.heap = nums
          heapq.heapify(self.heap)

          while len(self.heap) > k:
              heapq.heappop(self.heap)

      def add(self, val: int) -> int:
          heapq.heappush(self.heap, val)

          if len(self.heap) > self.k:
              heapq.heappop(self.heap)

          return self.heap[0]


_ Time Complexity:

  O(nlog(n) + mlog(k)) - The initial nlog(n) is for heapifying the input. The mlog(k) is for adding
  m elements to the heap and popping k elements.

_ Space Complexity:

  O(n) - We initially store n elements on the heap.