703. Kth Largest Element in a Stream (Python)

Related Topic

Heap.

Description

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Your KthLargest class will have a constructor which accepts an integer k and an integer array nums, which contains initial elements from the stream. For each call to the method KthLargest.add, return the element representing the kth largest element in the stream.

Sample I/O

Example 1

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

Note

You may assume that nums’ length ≥ k-1 and k ≥ 1.

Methodology

Use min heap and keep the min heap length to k. When adding a new number, compare it with the first element (min value) in heap. If it is greater than the min value then push the number and pop the mean value otherwise continue. The min heap will always keep the first kth largest values and the first element in min heap is always the kth largest value.

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        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:
        if len(self.heap) < self.k:
            heapq.heappush(self.heap,val)
        else:
            if val > self.heap[0]:
                heapq.heappushpop(self.heap,val)
        return self.heap[0]

BigO

Time complexity : heapify function here will take O(nlogn) where n is the size of input. push and pop will take O(logn) so sum up will be O(nlogn).