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
-
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 topk
is always at the root. -
Constructor: The constructor initializes the min-heap with the initial
nums
array. It iterates throughnums
, adding each element to the heap. If the heap size exceedsk
, the smallest element (root) is removed to maintain the heap size atk
. -
add() Method: The
add()
method adds a new number to the heap. Similar to the constructor, if the heap size exceedsk
after adding the new element, the smallest element is removed. Theadd()
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.
- Constructor: O(N log k), where N is the size of the input array
-
Space Complexity: O(k). The min-heap stores at most
k
elements.