703. Kth Largest Element in a Stream

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 <= 104
• 0 <= nums.length <= 104
• -104 <= nums[i] <= 104
• -104 <= val <= 104
• At most 104 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.

Approach 01:

class KthLargest {
public:
    // Constructor to initialize the KthLargest object with the given k and elements.
    KthLargest(int k, vector<int>& nums) : kthLargest(k) {
        for (const int num : nums)
            addElementToHeap(num);  // Add each element from the initial list to the heap.
    }

    // Function to add a new element and return the kth largest element.
    int add(int value) {
        addElementToHeap(value);
        return minHeap.top();  // The top of the min-heap is the kth largest element.
    }

private:
    const int kthLargest;  // The 'k' value indicating which largest element to maintain.
    priority_queue<int, vector<int>, greater<>> minHeap;  // Min-heap to keep track of the k largest elements.

    // Helper function to maintain the size of the heap to be at most k elements.
    void addElementToHeap(int value) {
        minHeap.push(value);  // Add the new element to the heap.
        if (minHeap.size() > kthLargest)
            minHeap.pop();  // Remove the smallest element if the heap exceeds size k.
    }
};

import heapq

class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.kthLargest = k
        self.minHeap = []
        for num in nums:
            self.addElementToHeap(num)  # Add each element from the initial list to the heap.

    def add(self, value: int) -> int:
        self.addElementToHeap(value)
        return self.minHeap[0]  # The top of the min-heap is the kth largest element.

    def addElementToHeap(self, value: int):
        heapq.heappush(self.minHeap, value)  # Add the new element to the heap.
        if len(self.minHeap) > self.kthLargest:
            heapq.heappop(self.minHeap)  # Remove the smallest element if the heap exceeds size k.

Time Complexity

• Constructor: KthLargest(int k, vector& nums):

  The constructor initializes the min-heap with the first k largest elements. Each insertion into the heap takes \( O(\log k) \) time, and since there are \( n \) elements in the input vector, the total time complexity for the constructor is \( O(n \log k) \).

• Function: int add(int value):

  Each call to the add function inserts a new element into the min-heap and possibly removes the smallest element if the heap size exceeds k. The insertion and removal operations both take \( O(\log k) \) time. Therefore, the time complexity for the add function is \( O(\log k) \).

Space Complexity

• Min-Heap:

  The min-heap stores at most k elements at any given time. Hence, the space complexity for storing the heap is \( O(k) \).

• Overall Space Complexity:

  In addition to the heap, the algorithm uses a constant amount of extra space for storing variables like kthLargest and minHeap. Therefore, the overall space complexity is \( O(k) \).