Kth Largest Element in a Stream

Easy Subject: Heap / Priority Queue

Time Complexity

O(N log K)

Space Complexity

O(K)

Problem Description

Problem Statement

Design a class to find the kth largest element in a stream of numbers. KthLargest(int k, int[] nums) initializes the object. int add(int val) appends val to the stream and returns the element representing the kth largest element.

Example 1:

Input: ["KthLargest", "add", "add", "add", "add", "add"]\n[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output: [null, 4, 5, 5, 8, 8]

Optimal Solution

Python

Approach: Min-Heap

Maintain a min-heap of size k. The top of the heap will always be the kth largest element. When adding, push the value to the heap. If the heap size exceeds k, pop the smallest element.

import heapq

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:
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]