Kth Largest Element in a Stream easy

Problem Statement

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 that accepts an integer k and an integer array nums, which contains initial numbers from the stream. The class should also have a method add, which accepts a new number and returns the kth largest element in the current 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]); creates a KthLargest object with k=3 and initial numbers [4, 5, 8, 2].
  • kthLargest.add(3); returns 4. The stream becomes [4, 5, 8, 2, 3]. The 3rd largest is 4.
  • kthLargest.add(5); returns 5. The stream becomes [4, 5, 8, 2, 3, 5]. The 3rd largest is 5.
  • kthLargest.add(10); returns 5. The stream becomes [4, 5, 8, 2, 3, 5, 10]. The 3rd largest is 5.
  • kthLargest.add(9); returns 8. The stream becomes [4, 5, 8, 2, 3, 5, 10, 9]. The 3rd largest is 8.
  • kthLargest.add(4); returns 8. The stream becomes [4, 5, 8, 2, 3, 5, 10, 9, 4]. The 3rd largest is 8.

Example 2

Input:

["KthLargest", "add", "add"]
[[1,[]], [-3],[-2]]

Output:

[null,-3,-2]

Explanation:

  • KthLargest kthLargest = new KthLargest(1, []); creates a KthLargest object with k=1 and an empty initial array.
  • kthLargest.add(-3); returns -3. The stream becomes [-3]. The 1st largest is -3.
  • kthLargest.add(-2); returns -2. The stream becomes [-3, -2]. The 1st largest is -2.

Steps

  1. Use a Min-Heap: We'll employ a min-heap data structure (priority queue in C++) to store the k largest elements encountered so far. A min-heap ensures that the smallest element among the top k is always at the root.

  2. Constructor: The constructor initializes the min-heap with the initial nums array. It iterates through nums, adding each element to the heap. If the heap size exceeds k, the smallest element (root) is removed to maintain the heap size at k.

  3. add() Method: The add() method adds a new number to the heap. Similar to the constructor, if the heap size exceeds k after adding the new element, the smallest element is removed. The add() method then returns the root of the heap (the kth largest element).

Explanation

The min-heap efficiently maintains the k largest elements seen so far. Adding a new element takes O(log k) time (heap insertion/deletion), which is efficient for large streams. Using a sorted array or repeatedly sorting the stream would be far less efficient.

Code

#include <queue>
#include <vector>

class KthLargest {
private:
    priority_queue<int, vector<int>, greater<int>> minHeap; // Min-heap
    int k;

public:
    KthLargest(int k, vector<int>& nums) : k(k) {
        for (int num : nums) {
            add(num);
        }
    }

    int add(int val) {
        minHeap.push(val);
        if (minHeap.size() > k) {
            minHeap.pop();
        }
        return minHeap.top();
    }
};

Complexity

  • Time Complexity:

    • Constructor: O(N log k), where N is the size of the input array nums. This is due to potentially adding N elements to the heap.
    • add() method: O(log k). Adding and removing from the heap take logarithmic time.
  • Space Complexity: O(k). The min-heap stores at most k elements.