703. Kth Largest Element in a Stream — LeetCode(Python)
I got you!
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 integerk
and the stream of integersnums
.int add(int val)
Appends the integerval
to the stream and returns the element representing thekth
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 toadd
. - It is guaranteed that there will be at least
k
elements in the array when you search for thekth
element.
Solution:
Explanation —
We can easily solve this problem using heaps.
First, in our constructor, we create two instance variables, self.nums and self.k and assign them the parameters that we were passed.
We then heapify the list that we were given.
Now, we pop elements from the heap until there are only k elements left. This is because we need to return the kth largest number, which is just the smallest element in a minheap of k elements.
Note that in Python, popping means removing the smallest element from the heap. Thus the larger elements remain unbothered.
In the add() method, we first push the new element into our heap, and then we pop an element to keep the heap length equal to k.
However, there is a caveat here. We should only pop from the heap if there are at least k+1 elements in it. We could possibly be given a list of length less than k to start with in the first place.
Finally, we return the root of the heap, which is the min element and also the kth largest element in the stream.
Time and Space Complexity —
Creation of the heap is a linear time operation if we use the heapify function. Push and pop operations require logarithmic time to run.
Also, we are using a linear amount of auxiliary space, since even though our heap eventually becomes of size k, we initially heapify the entire list nums.
Thus, if n is the length of the input array nums,
Time Complexity: O(nlogn + mlogk)
Space Complexity: O(n)