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)
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
104
calls will be made to add
.k
elements in the array when you search for the kth
element.Source
from bisect import bisect
class KthLargest():
"""naive solution"""
def __init__(self, k, nums):
self.k = k
self.array = sorted(nums)
def add(self, val):
"""Note: maintaining an array of only length k does not speed up the algo
and using a deque would be unefficient because of O(n) lookup
"""
idx = bisect(self.array, val) # ! bisect requires a sorted list
self.array.insert(idx, val)
k = -self.k
return self.array[-k]
import heapq
class KthLargest:
"""Better solution"""
def __init__(self, k, nums):
self.k = k
self.heap = nums
heapq.heapify(self.heap) # convert array into heap
while len(self.heap) > k:
heapq.heappop(self.heap)
def add(self, val):
"""Kth largest element --> maintain MinHeap of K elements and return the smallest
Kth smallest element --> maintain MaxHeap of K elements and return the largest
"""
if len(self.heap) < self.k:
heapq.heappush(self.heap, val)
elif val > self.heap[0]:
heapq.heappushpop(self.heap, val) # heapreplace() would work too
return self.heap[0]
# Note: difference between heappushpop(heap, val) and heapreplace(heap, val)
# heapreplace replace smallest value by val ==> if val < min(heap), we have now an even smaller min(heap)
# heappushpop would push val, and pop it out (since it is the min) --> heap would remained unchanged
kthLargest = KthLargest(3, [4, 5, 8, 2])
result = []
result.append(kthLargest.add(3))
result.append(kthLargest.add(5))
result.append(kthLargest.add(10))
result.append(kthLargest.add(9))
result.append(kthLargest.add(4))
print(result)
print(kthLargest.heap)
[4, 4, 4, 4, 4] [-4, -2, -3]
kthLargest = KthLargest(10,[-10,1,3,1,4,10,3,9,4,5,1])
result = []
result.append(kthLargest.add(3))
result.append(kthLargest.add(2))
result.append(kthLargest.add(3))
result.append(kthLargest.add(1))
result.append(kthLargest.add(2))
result.append(kthLargest.add(4))
result.append(kthLargest.add(5))
result.append(kthLargest.add(5))
result.append(kthLargest.add(6))
result.append(kthLargest.add(7))
result.append(kthLargest.add(7))
result.append(kthLargest.add(8))
result.append(kthLargest.add(2))
result.append(kthLargest.add(3))
result.append(kthLargest.add(1))
result.append(kthLargest.add(1))
result.append(kthLargest.add(1))
result.append(kthLargest.add(10))
result.append(kthLargest.add(11))
result.append(kthLargest.add(5))
result.append(kthLargest.add(6))
result.append(kthLargest.add(2))
result.append(kthLargest.add(4))
result.append(kthLargest.add(7))
result.append(kthLargest.add(8))
result.append(kthLargest.add(5))
result.append(kthLargest.add(6))
print(result)
print(kthLargest.heap)
[5, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2] [-2, -2, -1, -1, -1, -1, 10, -1, -1, -1]
Return the Kth smallest element instead of the Kth largest
import heapq
class KthLargest:
def __init__(self, k, nums):
"""only min heap in python --> -nums"""
self.k = k
self.heap = [-n for n in nums]
heapq.heapify(self.heap) # convert array into heap
while len(self.heap) > k:
heapq.heappop(self.heap)
def add(self, val):
"""Kth smallest element --> maintain MaxHeap of K elements and return the largest"""
if len(self.heap) < self.k:
heapq.heappush(self.heap, -val)
elif val > self.heap[0]:
heapq.heappushpop(self.heap, -val)
return -self.heap[0]