LeetCode 703: Kth Largest Element in a Stream (Python)

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

Solution:

from heapq import heapify, heappush, heappop, heappushpop

class KthLargest:

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.k = k
        self.k_heap = nums
        heapify(self.k_heap)
        while len(self.k_heap) > k:
            heappop(self.k_heap)

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if len(self.k_heap) < self.k:
            heappush(self.k_heap, val)
        else:
            heappushpop(self.k_heap, val)
        return self.k_heap[0]

Explanation:

On the initializer:

1- Create a min-heap with the stream of numbers we are given.

2- Pop from the min-heap until it has the same amount of elements as k . This means that the Kth largest element is on the top of the min-heap.

On add method:

1- If the amount of elements inside the heap is less than K , push the element to the heap. Else, push the element and then pop the top element.

2- Return the top element of the min-heap. It will be the Kth largest element.