Problem overview
You write a class KthLargest. Its constructor receives k and the numbers nums that have already arrived. Each call to add(value) adds 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 largest is 7, and so is the 3rd.
Take k = 3 and the starting numbers 4, 5, 8, 2. After add(3) the numbers are 2, 3, 4, 5, 8. The three largest are 8, 5 and 4, so the 3rd largest is 4. The five calls add(3), add(5), add(10), add(9), add(4) return 4, 5, 5, 8, 8.
Approach
The simple solution keeps every number in a list. On each add it appends the value, sorts the list in descending order and reads position k - 1. That works, but a sort costs O(n log n), where n is how many numbers have arrived, and it runs again on every call. The list also grows without limit.
The better idea starts from a question: which numbers can ever be the answer? New numbers only push the k-th largest up, never down. A number that is not among the k largest now will stay outside them forever, so it can be forgotten. It is enough to remember the k largest numbers.
Among those k numbers, the one you need is the smallest: the k-th largest is the smallest of the top k. A min-heap gives the smallest value in O(1) and lets you add or remove one in O(log size). So the class keeps a min-heap of at most k values:
- add the new value to the heap;
- if the heap now holds k + 1 values, remove the smallest, since it cannot be among the k largest;
- the top of the heap is the answer.
The heap never holds more than k + 1 values, so each add takes O(log k) time and the object uses O(k) space.
Step by step
- Declare the heap. A
val largestholds aPriorityQueue<Int>(). In Kotlin this is Java's priority queue, and with no comparator it is a min-heap: the smallest value is on top, which is the order the idea needs. - Feed the starting numbers. The numbers in
numsarrived before any call, but they count like any other. Aninitblock loops overnums. It runs when the object is created, after the properties above it, solargestalready exists. - Reuse
add. The loop callsadd(num)and ignores the result.addalready keeps the heap at k values, so the constructor does not repeat that logic. - Insert the new value into
largest. - Trim to k. If
largest.sizeis abovek, callpoll(), which removes and returns the smallest value. - Return the top with
peek(), which reads the smallest value without removing it. The constraints promise at least k numbers whenaddreturns, so the heap holds exactly k and its top is the k-th largest.
The solution in Kotlin
class KthLargest(private val k: Int, nums: IntArray) {
val largest = PriorityQueue<Int>()
init {
for (num in nums) {
add(num)
}
}
fun add(value: Int): Int {
largest.add(value)
if (largest.size > k) {
largest.poll()
}
return largest.peek()
}
}
Some Kotlin details that matter here:
kis a property,numsis not. The header saysprivate val k, sokstays available inaddafter the constructor ends.numshas noval, so it is a plain parameter, usable only in the property initializers andinitblocks.largest.addis not the class'sadd. Inside the class,add(num)calls the method being written, andlargest.add(value)calls the heap's method. The dot decides.poll()versuspeek(). Both look at the smallest value;poll()also removes it. Both come from Java'sPriorityQueue.IntArrayversusList. The constructor takes anIntArray, which is iterated withfor (num in nums)like any collection.
Seeing it run
This code creates the object, adds values and prints the heap's contents in sorted order next to each answer:
fun main() {
val kth = KthLargest(3, intArrayOf(4, 5, 8, 2))
println(kth.largest.sorted())
for (v in intArrayOf(3, 5, 10, 9, 4)) {
println("add($v) = ${kth.add(v)}, heap = ${kth.largest.sorted()}")
}
val single = KthLargest(1, intArrayOf())
println(listOf(-3, -2, -4, 0, 4).map { single.add(it) })
val ties = KthLargest(4, intArrayOf(7, 7, 7, 7, 8, 3))
println(ties.largest.sorted())
println(listOf(2, 10, 9, 9).map { ties.add(it) })
}
It prints:
[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]
[-3, -2, -2, 0, 4]
[7, 7, 7, 8]
[7, 7, 7, 8]
The heap holds the same three values as the list of the three largest numbers, and its smallest entry is always the answer. With k = 1 the heap holds one value, so every call returns the largest number so far. In the last example the 3 in the starting numbers was dropped as soon as the object was built, because four larger values were already there.
Complexity
- Time per
add: O(log k). The heap has at most k + 1 values, so one insertion and at most one removal cost O(log k) each. - Constructor: n starting numbers, each added in O(log k), so O(n log k).
- Space: O(k), because only the k largest values are kept.
Re-sorting everything costs O(n log n) per call and O(n) space, which is why the heap is the better design as the stream grows.
Edge cases and common mistakes
- Duplicates. Do not turn the heap into a set. With
k = 4and the numbers7, 7, 7, 7, 8, 3, the answer is 7 because four separate 7s can sit in the top four with the 8 above them. - Keeping a max-heap of everything. It would give the largest value quickly, not the k-th, and it would grow with the stream.
- Removing before adding. If you poll first and insert after, you throw away the smallest of the old top k even when the new value is smaller still, and the answer can be wrong. Insert first, then trim: the smallest of k + 1 values is the right one to drop.
- Trimming at the wrong size. The heap must be cut back when it holds more than
kvalues. Cutting at another size leaves it one value off, andpeek()then returns the wrong rank. - Few starting numbers. The heap may hold fewer than k values after the constructor, for example
KthLargest(1, intArrayOf()). That is fine, because the constraints promise at least k values by the timeaddreturns.
