Problem overview
Numbers arrive one at a time, and after each one you report the k-th largest of everything seen so far. Equal numbers count separately: in [8, 7, 7, 5] the 2nd largest is 7 and so is the 3rd.
The struct KthLargest has two methods. KthLargest::new(k, nums) builds the object from k and the numbers that have already arrived. add(val) takes in one more number and returns the k-th largest so far. The constraints promise that at least k numbers have arrived whenever add returns.
With k = 3 and nums = [4, 5, 8, 2], calling add(3) leaves the numbers [2, 3, 4, 5, 8]. The three largest are 8, 5 and 4, so the answer is 4.
Approach: keep only the k largest
The simple way is to store every number, sort the list on each add and read the element at position k from the top. That costs O(n log n) per call, where n is how many numbers have arrived, and the object keeps all n of them. The statement asks for O(log k) per call and O(k) memory.
The key observation is that a number outside the k largest can never matter again. Numbers only arrive, never leave, so the k-th largest can only stay the same or grow. A number that is already below it falls further behind with every arrival. So the object needs to remember only the k largest values.
Among those k values, the one the object needs is the smallest, because it is the k-th largest overall. A min-heap gives the smallest value at the top in O(1) and supports insertion and removal of the top in O(log k). Each add pushes the new value, and if the heap now holds k + 1 values, pops the smallest. The top is then the answer.
Rust's BinaryHeap is a max-heap, so the lesson wraps each value in Reverse, from std::cmp, which flips the order. With plain integers 5, 1 and 9 pushed, peek() gives Some(9); with each wrapped in Reverse, it gives Some(Reverse(1)).
Step by step
- Store
kin the struct as ausize, so it can be compared with the heap'slen()directly. - Add the heap field,
BinaryHeap<Reverse<i32>>: a min-heap that holds the k largest values, smallest on top. - In
new, build the struct withk as usizeand an empty heap. It is declaredmutbecause the next steps change it. - Loop over
nums: the starting numbers must pass through the same filter as every later one. - Call
addfor each number. It already keeps the heap at k values, so the constructor reuses it. The returned values are ignored. - Return the struct as the last expression, without a semicolon.
- In
add, pushReverse(val)into the heap. - If the heap holds more than
self.kvalues, pop its smallest. That value cannot be among the k largest now or later. - Return the top:
peek()gives anOption,unwrap()takes theReverseout, and.0is the value inside it.
The solution in Rust
use std::cmp::Reverse;
use std::collections::BinaryHeap;
struct KthLargest {
k: usize,
largest: BinaryHeap<Reverse<i32>>,
}
impl KthLargest {
fn new(k: i32, nums: Vec<i32>) -> KthLargest {
let mut kth_largest = KthLargest { k: k as usize, largest: BinaryHeap::new() };
for num in nums {
kth_largest.add(num);
}
kth_largest
}
fn add(&mut self, val: i32) -> i32 {
self.largest.push(Reverse(val));
if self.largest.len() > self.k {
self.largest.pop();
}
self.largest.peek().unwrap().0
}
}
Using it with k = 3 and nums = [4, 5, 8, 2] gives this. The heap column lists the heap's contents in sorted order.
start heap [4, 5, 8]
add(3) -> 4 heap [4, 5, 8]
add(5) -> 5 heap [5, 5, 8]
add(10) -> 5 heap [5, 8, 10]
add(9) -> 8 heap [8, 9, 10]
add(4) -> 8 heap [8, 9, 10]
The constructor pushed 4, 5, 8 and 2, and the 2 was popped as soon as the fourth value arrived. After add(3) the 3 is pushed and immediately popped again, so the heap does not change. The 5 pushes out the 4, the 10 pushes out one of the two 5s, and the final add(4) is below the top, so it leaves and nothing changes.
The Rust details worth a closer look:
peek()returns anOption. It isNonefor an empty heap, andunwrap()would panic then. Hereaddhas just pushed a value, so the heap is never empty at that point.- Reading the wrapper.
Reverse(val)is a tuple struct, and.0reads its single field, which is howpeek().unwrap().0gets the integer back out. pop()as a statement. It returns anOptionwith the removed value; the code discards it. The final line has no semicolon, so it is the method's return value.&mut self.addchanges the heap, so it borrows the struct mutably, and the constructor's local variable must bemut.- Casts.
k as usizeconverts thei32parameter once, in the constructor. The constraints keepkat 1 or more, so the cast is safe.
Complexity
Time: O(log k) per add. The heap never holds more than k + 1 values, so one push and at most one pop each cost O(log k), and peek is O(1). The constructor makes one add call per starting number, so it takes O(n log k) for n numbers.
Space: O(k). The heap holds at most k + 1 values at any moment, and exactly k between calls once enough numbers have arrived.
Edge cases and common mistakes
- Duplicates. They count separately. With
k = 4andnums = [7, 7, 7, 7, 8, 3],add(2)returns 7, and so doadd(10)andadd(9). A secondadd(9)pushes out the last 7 and returns 8. Putting the numbers in a set would break this. - An empty
nums. Withk = 1and no starting numbers,add(-3),add(-2),add(-4),add(0)andadd(4)return -3, -2, -2, 0 and 4: the largest so far. - Using a max-heap. Without
Reverse, the top is the largest value, not the k-th largest, and popping removes the wrong end. - Popping before pushing. The new value has to enter the heap first, since it may belong among the k largest. Popping first would drop a value that should stay.
- Fewer than k numbers. The constraints rule it out for the answer you return. If it did happen, the heap would hold fewer than k values and the top would be the smallest of them, not a k-th largest. For example, with
k = 3andnums = [5],add(1)returns 1. - Sorting on every call. It works but costs
O(n log n)per call and keeps all the numbers.
Testing
k = 3,nums = [4, 5, 8, 2], thenadd(3),add(5),add(10),add(9),add(4)returns 4, 5, 5, 8, 8.k = 1,nums = [], thenadd(-3),add(-2),add(-4),add(0),add(4)returns -3, -2, -2, 0, 4.k = 4,nums = [7, 7, 7, 7, 8, 3], thenadd(2),add(10),add(9),add(9)returns 7, 7, 7, 8.
