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(val)
, which 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
Example 2
Input
["KthLargest", "add", "add"]
[[2, [0]], [-1], [1]]
Output
[null, 0, 1]
Explanation
KthLargest kthLargest = new KthLargest(2, [0]);
kthLargest.add(-1); // return 0
kthLargest.add(1); // return 1
Steps
-
Data Structure: Use a min-heap (priority queue) to store the k largest elements encountered so far. A min-heap ensures that the smallest of the k largest elements is always at the root.
-
Constructor: Initialize the min-heap with the initial
nums
array. If the heap size exceedsk
, remove the smallest element (root) to maintain the heap size atk
. -
add(val)
method: Add the new elementval
to the heap. If the heap size exceedsk
, remove the smallest element. Return the root element (smallest of the k largest), which is the kth largest element.
Explanation
The key to solving this problem efficiently is using a min-heap. A min-heap keeps track of the smallest k
elements. Adding a new element, if it's larger than the smallest of the current k
elements, pushes the smallest element out and makes room for the new larger element. Therefore, the root will always represent the kth largest element.
Code
import { MinPriorityQueue } from '@datastructures-js/priority-queue';
class KthLargest {
k: number;
minHeap: MinPriorityQueue<number>;
constructor(k: number, nums: number[]) {
this.k = k;
this.minHeap = new MinPriorityQueue({ priority: (x) => x });
for (const num of nums) {
this.add(num);
}
}
add(val: number): number {
this.minHeap.enqueue(val);
if (this.minHeap.size() > this.k) {
this.minHeap.dequeue();
}
return this.minHeap.front()!.element;
}
}
Note: This code uses the @datastructures-js/priority-queue
package for the min-heap implementation. You'll need to install it using npm install @datastructures-js/priority-queue
. If you prefer a different min-heap implementation (like a custom one using a binary heap), you can replace this with your own.
Complexity
-
Time Complexity:
- Constructor: O(N log k), where N is the length of
nums
. Adding each element to the heap takes O(log k) time. add(val)
: O(log k). Adding and removing elements from the heap takes O(log k) time.
- Constructor: O(N log k), where N is the length of
-
Space Complexity: O(k), the space used by the min-heap.