Problem overview
The method kSmallest receives an Array[Int] of positive integers and a number k. It returns an array holding the k smallest values of nums, in any order. Duplicates count separately: for nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4 the four smallest are 1, 2, 2 and 3, because the array has two 2s.
Approach: keep only the k best, with the worst on top
The easy solution sorts the whole array and takes the first k values, which costs O(n log n) time and needs room for the sorted copy. This lesson sets a tighter budget, O(n log k) time and O(k) extra space, which matters when k is much smaller than n.
The idea is to hold the k smallest values seen so far in a max heap. A max heap gives instant access to its largest element, and here the largest element of the kept values is the one that is first to be thrown out. For each new value:
- if fewer than
kvalues are kept, add it; - otherwise compare it with the largest kept value. If it is smaller, that largest value can't be among the
ksmallest anymore, so remove it and add the new one. If it is not smaller, ignore it.
Each heap operation takes O(log k) because the heap never holds more than k values, and there are at most two per element. That gives O(n log k).
For the example, the first four values 5, 9, 3, 6 fill the heap. The next value, 2, is smaller than the top (9), so 9 leaves and 2 enters, giving 5, 3, 6, 2. Then 1 replaces 6, 3 replaces 5, 2 replaces 3 (the heap now holds 3, 1, 2, 2), and the remaining values 7 and 5 are not smaller than the top (3), so they are ignored. The heap ends with 1, 2, 2 and 3.
Step by step
- Create the heap. Declare
smallestKNumsas avalholding an emptyscala.collection.mutable.PriorityQueue[Int](). Itsdequeuereturns the greatest element first, so it already behaves as a max heap. - Loop over the values with
for (value <- nums). - Fill the heap first. If
smallestKNums.size < k, callenqueue(value). - Compare with the top. Otherwise, in an
else if, testvalue < smallestKNums.head.headreads the greatest element without removing it. - Swap. Inside that branch call
dequeue()to remove the greatest value, thenenqueue(value). The heap stays at exactlykvalues. - Copy the heap into an array. After the loop, create
resultasnew Array[Int](k). In a loop over0 until k, storesmallestKNums.dequeue()intoresult(i). - Return
resultas the last expression. The values come out from largest to smallest, which is accepted because the order is free.
The solution in Scala
class Solution {
def kSmallest(nums: Array[Int], k: Int): Array[Int] = {
val smallestKNums = scala.collection.mutable.PriorityQueue[Int]()
for (value <- nums) {
if (smallestKNums.size < k) {
smallestKNums.enqueue(value)
} else if (value < smallestKNums.head) {
smallestKNums.dequeue()
smallestKNums.enqueue(value)
}
}
val result = new Array[Int](k)
for (i <- 0 until k) {
result(i) = smallestKNums.dequeue()
}
result
}
}
On the example, the method returns [3, 2, 2, 1]. The Scala details worth a closer look:
- A max heap by default.
PriorityQueue[Int]()uses the implicit ordering ofIntand puts the greatest element on top. For a min heap, pass the reversed ordering explicitly:PriorityQueue[Int]()(Ordering[Int].reverse). - Printing a heap doesn't sort it. After enqueueing 5, 1, 4, 2, 3,
println(pq)showsPriorityQueue(5, 3, 4, 1, 2), the heap's internal layout. Onlyheadanddequeueare guaranteed to follow the priority. That is why the code drains the queue withdequeue()instead of copying it. headversusdequeue().headonly looks;dequeue()removes and returns. Both throwNoSuchElementExceptionon an empty queue.new Array[Int](k)creates an array ofkzeros, and the loop overwrites every slot. Arrays are indexed and updated with parentheses:result(i) = ....0 until kexcludesk, so the loop writes indices 0 tok - 1.- Method ends with
result. Without areturn, the last expression is the method's value.
Complexity
Time: O(n log k). Each of the n values costs one comparison with head and, at most, one dequeue and one enqueue on a heap of size k, each O(log k). Emptying the heap at the end costs O(k log k), which is no more than O(n log k) when k <= n.
Space: O(k). The heap and the result array each hold k values.
Edge cases and common mistakes
kequal to the array length. Every value goes into the heap and nothing is replaced. The result is the whole array, largest first.k = 1. The heap keeps one value, the smallest so far.[5, 9, 3]returns[3].k = 0orklarger than the array. Both are outside what the method handles. Withk = 0the heap stays empty, so readingheadthrowsNoSuchElementException. Withk = 5and three values, the final loop dequeues from an empty heap and throws the same exception. Assume1 <= k <= nums.length.- Using a min heap. A min heap would put the smallest value on top, which is the one you want to keep. The heap must expose the value to discard, the largest.
- Duplicates. They are kept. Using a
Setwould silently drop the second 2 in the example. - Expecting sorted output. The result is largest first. Sort it if you need ascending order.
Testing
nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5],k = 4returns[3, 2, 2, 1], which holds the same values as[1, 2, 2, 3].nums = [5, 9, 3],k = 3returns[9, 5, 3].nums = [5, 9, 3],k = 1returns[3].
