Problem overview
Numbers arrive one at a time, and after each one you must 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 lesson asks for a class KthLargest(k: Int, nums: Array[Int]) with one method, add(value: Int): Int. The constructor receives the numbers that arrived before any call, and each add stores a new number and returns the current k-th largest. The statement promises that at least k numbers have arrived whenever add returns.
Example: KthLargest(3, Array(4, 5, 8, 2)) followed by add(3), add(5), add(10), add(9) and add(4) returns 4, 5, 5, 8 and 8. After add(3) the numbers are 2, 3, 4, 5, 8, the three largest are 8, 5 and 4, and the 3rd largest is 4.
Approach
The simple way and its cost
Keep every number in a buffer, sort it on each add, and read the element k places from the top. It works, but a sort costs O(n log n), where n is how many numbers have arrived, and the cost grows with every call. The buffer also grows without limit. The lesson asks for O(log k) time per add and O(k) memory.
Only the k largest matter
A number that is not among the k largest can never become the answer. New numbers only push the k-th largest up, never down, so a number that falls behind stays behind. That means the object can forget everything except the k largest values.
Among those k values, the answer is the smallest one. A min-heap gives the smallest of its values in O(1) and removes it in O(log k), which is exactly what is needed:
- push the new value into the heap;
- if the heap now holds
k + 1values, remove the smallest, which cannot be among the k largest; - the smallest value that is left is the k-th largest.
The heap never holds more than k + 1 values, so each push and each removal costs O(log k).
A min-heap in Scala
Scala's scala.collection.mutable.PriorityQueue is a max-heap: head is the largest value. To get a min-heap, hand it the reversed ordering in a second parameter list: PriorityQueue[Int]()(Ordering[Int].reverse). The first pair of parentheses is the (empty) list of starting elements, and the second one holds the ordering. With the reversed ordering, "largest" means "smallest number", so head is the smallest value.
Step by step
- Declare the heap. A
valnamedlargest, aPriorityQueue[Int]with the reversed ordering, holds the k largest values so far, smallest on top. - Loop over
nums. Statements in a class body run when the object is created, in order, solargestalready exists when the loop starts. - Call
add(num)for each starting number. The starting numbers must go through the same filter as every later one, andaddalready keeps the heap at k values. The constructor ignores the returned value, because nobody has asked for it yet. - In
add, enqueue the value. It might be one of the k largest; the next step finds out. - Trim the heap. If
largest.sizeis greater thank, calldequeue()to drop the smallest value. - Return
largest.head. It is the smallest of the k largest values, which is the k-th largest.
The solution in Scala
class KthLargest(k: Int, nums: Array[Int]) {
val largest = PriorityQueue[Int]()(Ordering[Int].reverse)
for (num <- nums) {
add(num)
}
def add(value: Int): Int = {
largest.enqueue(value)
if (largest.size > k) {
largest.dequeue()
}
return largest.head
}
}
Some Scala details:
- Class parameters are visible in the whole class.
kandnumshave noval, yet the methods and the constructor loop can use them. Scala keeps a parameter as a private field when a method needs it. - The loop is the constructor. A
forloop in the class body is not inside any method, so it runs once, whennew KthLargest(...)is evaluated. It callsadd, which is defined further down; that is fine because methods exist before the body runs. The order of fields matters, though:largestis declared above the loop becauseadduses it. enqueueanddequeue().enqueueadds a value, anddequeue()removes and returns the top. The call todequeue()inaddthrows its result away; only the removal matters.headreads without removing. On an empty queue,headthrowsNoSuchElementException.- The lesson's code uses
PriorityQueuewithout a prefix. In a file of your own, addimport scala.collection.mutable.PriorityQueuefirst.
Complexity
Time per add: O(log k). One enqueue and at most one dequeue(), each on a heap of at most k + 1 values. Reading head is O(1). The constructor calls add once per starting number, so it costs O(m log k) for m numbers.
Space: O(k). The heap holds at most k + 1 values at any moment and k between calls, however many numbers arrive.
Edge cases and common mistakes
- Using a max-heap. Without
Ordering[Int].reversethe top is the largest value, which is the 1st largest, not the k-th. - Starting numbers. They count like any other number. Skipping the loop gives wrong answers whenever
numsis not empty. - Empty
nums.KthLargest(1, Array())starts with an empty heap. Because at least k numbers have arrived wheneveraddreturns,headis safe. - Duplicates. They are separate values in the heap. With
k = 4and starting numbers7, 7, 7, 7, 8, 3, the heap keeps four values, 7, 7, 7 and 8, and afteradd(2)the answer is 7, the smallest of them. - Trimming with
>=instead of>. The test issize > k. With>=, the heap shrinks tok - 1values andheadbecomes the (k-1)-th largest: forKthLargest(3, Array(4, 5, 8, 2)),add(3)returns 5 instead of 4. Withk = 1the heap is empty after trimming, andheadthrowsNoSuchElementException. - Fewer than k numbers. The heap then holds everything seen, and
headis the smallest of them, not a k-th largest. ForKthLargest(3, Array(5)),add(1)returns 1. The statement rules this case out.
Testing
KthLargest(3, Array(4, 5, 8, 2))withadd(3),add(5),add(10),add(9),add(4)returns4 5 5 8 8.KthLargest(1, Array())withadd(-3),add(-2),add(-4),add(0),add(4)returns-3 -2 -2 0 4.KthLargest(4, Array(7, 7, 7, 7, 8, 3))withadd(2),add(10),add(9),add(9)returns7 7 7 8.KthLargest(2, Array())withadd(9)returns9, which is outside the promise: only one number has arrived, so 9 is not a 2nd largest.
