The problem
Design a class that is given a number k and a starting list of numbers nums. Each call to add(val) adds one more number and returns the k-th largest of all the numbers so far (counting repeats).
Examples
01
- Input
["KthLargest", "add", "add", "add", "add", "add"] [[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
- Output
[null, 4, 5, 5, 8, 8]
02
- Input
["KthLargest", "add", "add", "add", "add"] [[4, [7, 7, 7, 7, 8, 3]], [2], [10], [9], [9]]
- Output
[null, 7, 7, 7, 8]
Constraints
- 0 ≤ nums.length ≤ 10⁴
- 1 ≤ k ≤ nums.length + 1
- −10⁴ ≤ nums[i], val ≤ 10⁴
- At most 10⁴ calls to
add.
The idea
A min-heap is a binary tree kept so every parent is no larger than its children: the smallest item is always at the top, and adding or removing an item takes log n steps, a swap per level.
Keep only the k largest numbers seen, in a min-heap. The smallest of those k — the top — is exactly the k-th largest. When a new number arrives, add it; if the heap now holds k + 1, remove the top, which can no longer be among the k largest.
- Time
- O(log k) per add
- Space
- O(k)
Solution · every language run against every case
class KthLargest: def __init__(self, k: int, nums: List[int]): # A min-heap of the k largest so far: its top, the smallest of them, is the k-th largest. self.k = k self.heap = [] for x in nums: self.add(x) def add(self, val: int) -> int: heappush(self.heap, val) if len(self.heap) > self.k: heappop(self.heap) # no longer among the k largest return self.heap[0]