View on NeetCode
View on LeetCode

Problem

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.

Implement KthLargest class:

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

Example 1:

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

Explanation:
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

Constraints:

  • 1 <= k <= 10^4
  • 0 <= nums.length <= 10^4
  • -10^4 <= nums[i] <= 10^4
  • -10^4 <= val <= 10^4
  • At most 10^4 calls will be made to add.
  • It is guaranteed that there will be at least k elements in the array when you search for the kth element.

Solution

Approach: Min Heap of Size K

The key insight is to maintain a min heap of size k. The root of the heap is always the kth largest element.

Implementation

import heapq

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

        # Keep only k largest elements
        while len(self.heap) > k:
            heapq.heappop(self.heap)

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

        # If heap size exceeds k, remove smallest
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)

        return self.heap[0]

Complexity Analysis

  • Time Complexity:
    • __init__: O(n log k), where n is the initial array size
    • add: O(log k) for heap operations
  • Space Complexity: O(k) for the heap storing k elements.

Key Insights

  1. Min Heap Property: In a min heap of size k, the smallest element (root) is the kth largest overall.

  2. Size Constraint: We maintain exactly k elements. When adding causes size > k, we remove the smallest to maintain the k largest.

  3. Efficient Updates: Using a heap allows O(log k) additions instead of O(n log n) sorting each time.

  4. Root is Answer: The heap root always gives us the kth largest element in O(1) time.

  5. Why Min Heap: A min heap keeps the smallest of the k largest elements at the root, which is exactly the kth largest element.