Problem overview
The method kSmallest receives an Array[Int] named nums of positive integers and an Int named k. It returns an array with the k smallest values, in any order. For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the sorted array starts 1, 2, 2, 3, so the answer is those four values, for example [1, 2, 2, 3] or [2, 1, 2, 3].
The statement allows O(n log n) time. Sorting a copy and taking the first k elements would meet that, but this lesson shows another route: a binary search over the values themselves.
Approach: find the k-th smallest value, then collect
The plan has two phases:
- Find the value of the k-th smallest element, call it
kthNum. - Collect every element smaller than
kthNum, then fill the remaining slots with copies ofkthNum.
The second phase works because of duplicates. All elements strictly smaller than kthNum belong in the answer, and there are fewer than k of them. The rest of the answer must equal kthNum, since that value is the k-th in sorted order. In the example kthNum is 3. The elements below 3 are 1, 2 and 2, which is three values, so exactly one copy of 3 completes the answer.
The first phase needs a test that can be searched. Define smaller(v) as how many elements of nums are strictly less than v. This count never decreases as v grows. The k-th smallest value is the largest v for which smaller(v) < k: at v = kthNum fewer than k elements are smaller, and for any larger v the k-th element itself is also smaller, which makes the count at least k. So the search can run over the range of values, from the minimum to the maximum of nums, looking for the largest v that passes the test.
Step by step
- Write
howManyLess(nums, value). It starts avar counterat 0, loops overnumswithnum, and adds one whenevervalue > num. It ends withcounter. - Write
kSmallestNum(nums, k). First setleftandrighttonums(0), then loop overnumsto turn them into the minimum and the maximum usingmath.minandmath.max. Startresultatleft. - While
left <= right, takevalue = left + (right - left) / 2andsmaller = howManyLess(nums, value). - If
smaller < k, this value is still a candidate for the k-th smallest, so save it inresultand try larger ones withleft = value + 1. Otherwise it is too large, so setright = value - 1. - Return
result. - In
kSmallest, storekthNum, createresultasnew Array[Int](k)and keep the next free slot invar i. - Loop over
nums; everyvaluebelowkthNumgoes intoresult(i), andimoves on by one. - While
i < k, writekthNumintoresult(i)and move on. Returnresult.
The solution in Scala
class Solution {
def kSmallestNum(nums: Array[Int], k: Int): Int = {
var left = nums(0)
var right = nums(0)
for (num <- nums) {
left = math.min(left, num)
right = math.max(right, num)
}
var result = left
while (left <= right) {
val value = left + (right - left) / 2
val smaller = howManyLess(nums, value)
if (smaller < k) {
result = value
left = value + 1
} else {
right = value - 1
}
}
result
}
def howManyLess(nums: Array[Int], value: Int): Int = {
var counter = 0
for (num <- nums) {
if (value > num) {
counter += 1
}
}
counter
}
def kSmallest(nums: Array[Int], k: Int): Array[Int] = {
val kthNum = kSmallestNum(nums, k)
val result = new Array[Int](k)
var i = 0
for (value <- nums) {
if (value < kthNum) {
result(i) = value
i += 1
}
}
while (i < k) {
result(i) = kthNum
i += 1
}
result
}
}
Scala details
The parameter is called value and not val, because val is a keyword in Scala.
new Array[Int](k) creates an array of k elements, all set to 0. Arrays are mutable and have a fixed size, so you fill slots with result(i) = ..., using parentheses for both reading and writing.
Scala has no i++, so storing a value and moving the position are two statements: result(i) = value and then i += 1.
The methods end with a bare name (result, counter), which is the method's value. Inside the loop, value, smaller are vals because each iteration makes new ones, while left, right, result, i and counter change and are vars.
Methods inside a class can call each other regardless of the order in which they are written, which is why kSmallestNum may call howManyLess even though it is defined first.
Complexity
The binary search runs over the range from the smallest to the largest value, so it makes about log(max - min) iterations. Each iteration calls howManyLess, which scans the array in O(n). The total is O(n log(max - min)). For positive 32-bit integers the range is below 2 to the power 31, so the logarithm is at most 31, which fits comfortably within the O(n log n) allowed by the statement for any realistic array. The final pass is O(n).
The extra space is the result array of size k, plus a few variables. Apart from the output, it is O(1).
Edge cases and common mistakes
- Duplicates of the k-th value. With
[4, 4, 4, 4]andk = 2, nothing is smaller than 4, so the first loop stores nothing and the padding loop writes two 4s. Without the padding step the two slots would stay 0. - k equal to 1. The k-th smallest is the minimum, so for the example array the answer is
[1]. - k equal to the length. For the example array and
k = 10, the answer contains every element, with the maximum 9 added last by the padding loop. The order is not sorted, which is allowed. - A single element.
[7]withk = 1gives[7]. - Using "less than or equal" in the count. The test
smaller < krelies on a strict count. Counting values less than or equal tovaluechanges which condition identifies the k-th value, and the padding logic no longer lines up. - Recording the result. The loop ends with
leftpastright, soresultmust keep the last candidate that passed the test. - k out of range. The method assumes
1 <= k <= nums.length. Withk = 0it returns an empty array, and akabove the length gives a meaningless result.
Testing
[5, 9, 3, 6, 2, 1, 3, 2, 7, 5]withk = 4gives[2, 1, 2, 3];kSmallestNumreturns 3 for it, andhowManyLesswith value 3 returns 3.- The same array with
k = 1gives[1]. [4, 4, 4, 4]withk = 2gives[4, 4].[7]withk = 1gives[7].
