Problem overview
The class KthLargest is built from an integer k and a list nums of numbers that have already arrived. Each call to add(val) takes in one more number and returns the k-th largest of everything seen so far. Equal numbers count separately: in [8, 7, 7, 5] the 2nd and the 3rd largest are both 7.
Take k = 3 and nums = [4, 5, 8, 2]. Calling add(3) makes the numbers [2, 3, 4, 5, 8]. The three largest are 8, 5 and 4, so the answer is 4. Then add(5) gives 5, add(10) gives 5, add(9) gives 8 and add(4) gives 8. The tests promise that at least k numbers have arrived whenever add returns, so the k-th largest always exists.
Approach: keep only the k largest
The simple approach keeps every number in a list. On each add it appends the new number, sorts the list and reads the element at position -k. That costs O(n log n) per call, where n is how many numbers have arrived, and the object stores all of them.
The better idea starts from one observation: a number that is not among the k largest can never become the answer. Numbers only arrive, never leave, so the k-th largest can only stay the same or go up. Anything smaller than it would only fall further behind. So the object keeps exactly k values, the k largest seen so far.
Among those k values, the answer is the smallest one: it has k-1 larger values above it, which makes it the k-th largest. What the object needs is a structure that always gives the smallest value quickly and lets a new value in cheaply. That structure is a min-heap, and Python's heapq module turns an ordinary list into one. The smallest element sits at index 0, and heappush and heappop each take O(log k) time on a heap of k values.
The rule for a new value is: push it, and if the heap now holds more than k values, pop the smallest. If the new value was small, it is the one that leaves. If it was big, it pushes out the old k-th largest, and the next smallest becomes the answer.
Step by step
- In
__init__, storekinself.k, becauseaddneeds it on every call. - Set
self.largestto an empty list. It will be the heap. Values go in as they are: a max-heap would need negated values, but here the smallest of the k largest is what we want on top, which is exactly whatheapqgives. - The numbers in
numsarrived before any call, but they count like all the others. Loop over them in the constructor. - Pass each one to
self.add. The constructor reuses the method, so the rule "keep only k values" is written once. The returned value is ignored, because nobody has asked yet. - In
add, pushvalontoself.largestwithheapq.heappush. - If the heap now holds more than
self.kvalues, remove the smallest withheapq.heappop. The heap never holds more than k + 1 values, which is why a push and a pop each costO(log k). - Return
self.largest[0], the smallest of the k largest values.
The solution in Python
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.largest = []
for num in nums:
self.add(num)
def add(self, val: int) -> int:
heapq.heappush(self.largest, val)
if len(self.largest) > self.k:
heapq.heappop(self.largest)
return self.largest[0]
Here is the class in use, with k = 3 and the starting numbers [4, 5, 8, 2]. Right after the constructor the heap holds the three largest, 4, 5 and 8, so the 2 was pushed out:
k = KthLargest(3, [4, 5, 8, 2])
print(k.largest)
print(k.add(3), k.largest)
print(k.add(5), k.largest)
print(k.add(10), k.largest)
This prints:
[4, 5, 8]
4 [4, 5, 8]
5 [5, 5, 8]
5 [5, 10, 8]
Adding 3 changes nothing: it was pushed and immediately popped, being the smallest. Adding 5 pushes the 4 out, leaving 5, 5 and 8. Adding 10 pushes out one of the 5s, leaving 5, 8 and 10, whose smallest is 5. The list is printed in the heap's internal order, so only index 0 is guaranteed to be the smallest; the other positions are not sorted.
A few details worth knowing:
heapqworks on a plain list and has no heap class, so every call passes the list as its first argument.heapqis always a min-heap. For the k smallest you would store negated values instead.- The lesson's editor already has
heapqandListavailable. In a file of your own, addimport heapqandfrom typing import List.
Complexity
Time: each add does one push and at most one pop on a heap of at most k + 1 values, so it takes O(log k). The constructor calls add once per number in nums, so it takes O(m log k) for m starting numbers. The sorting approach pays O(n log n) on every call.
Space: O(k). The object never stores more than k + 1 values, however many numbers arrive.
Edge cases and common mistakes
- Fewer than k starting numbers. With
k = 1andnums = []the heap starts empty. The firstaddpushes one value, nothing is popped because 1 is not more than 1, and that value is returned. If the heap holds fewer thankvalues, nothing is popped. The tests only calladdonce at least k numbers exist, soself.largest[0]is always the k-th largest. - Duplicates. They are separate values in the heap. With
k = 4and[7, 7, 7, 7, 8, 3]the heap holds 7, 7, 7 and 8 after the constructor, so the 4th largest is 7. Afteradd(2)the answer is still 7, because the 2 is smaller and is popped at once. - Negative numbers. Nothing special: the heap compares them like any other integers.
- Returning the wrong end. The top of a min-heap is the smallest value in the heap. That is the answer here only because the heap was trimmed to k values. Popping before pushing, or trimming to k - 1, shifts the answer by one place.
Testing
These inputs, with the answers printed by the solution above, are worth trying:
k1 = KthLargest(1, [])
print([k1.add(v) for v in (-3, -2, -4, 0, 4)])
k4 = KthLargest(4, [7, 7, 7, 7, 8, 3])
print([k4.add(v) for v in (2, 10, 9, 9)])
[-3, -2, -2, 0, 4]
[7, 7, 7, 8]
With k = 1 the answer is simply the largest number so far. In the second run, the second 9 pushes the last 7 out, so the four largest become 10, 9, 9 and 8 and the answer is 8.
