← Back to Index

Kth largest element in a stream of running integers

Implement a class KthLargest to find the kth largest number in a stream. It should have the following methods:

KthLargest(int k, int [] nums) Initializes the object with the integer k and the initial stream of numbers in nums int add(int val) Appends the integer val to the stream and returns the kth largest element in the stream.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

Input: [KthLargest(3, [1, 2, 3, 4]), add(5), add(2), add(7)]
Output: [null, 3, 3, 4]
Explanation: initial stream = [1, 2, 3, 4], k = 3.
add(5): stream = [1, 2, 3, 4, 5] -> returns 3
add(2): stream = [1, 2, 2, 3, 4, 5] -> returns 3
add(7): stream = [1, 2, 2, 3, 4, 5, 7] -> returns 4


Input: [KthLargest(2, [5, 5, 5, 5], add(2), add(6), add(60)]
Output: [null, 5, 5, 6]
Explanation: initial stream = [5, 5, 5, 5], k = 2.
add(2): stream = [5, 5, 5, 5, 2] -> returns 5
add(6): stream = [5, 5, 5, 5, 2, 6] -> returns 5
add(60): stream = [5, 5, 5, 5, 2, 6, 60] -> returns 6

Approach 1 (Min heap)

• while initializing , build a min heap of size k
• on add operation , add the element to min_heap, ensure heap size and return the kth element
• create a helper balance which will ensures size of heap cannot exceed k

import heapq

class KthLargest:
    def __init__(self,k,nums):
        self.min_heap = []
        self.k = k
                
        for num in nums:
            heapq.heappush(self.min_heap,num)

        self.balance()

        

    def add(self,num):
        heapq.heappush(self.min_heap, num)
        self.balance()
    
        # Return the kth largest (smallest in min-heap) without removing it
        return self.min_heap[0]

    def balance(self):
        while len(self.min_heap) > self.k:
            heapq.heappop(self.min_heap)

        
obj1 = KthLargest(3,[1, 2, 3, 4])

print(obj1.add(5))
print(obj1.add(2))
print(obj1.add(7))

obj1 = KthLargest(2,[5, 5, 5, 5])

print(obj1.add(2))
print(obj1.add(6))
print(obj1.add(60))

3
3
4
5
5
6

Complexity

• N= len(nums)

• O(N log(k)) for constructor , O(log K) for add

• O(K) for storing min_heap