Problem overview
The function kSmallest receives a slice of positive integers nums and a number k. It returns the k smallest values of nums, in any order. Duplicates count separately: if the value 2 appears twice and both are among the smallest, both appear in the answer.
For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the sorted values are 1, 2, 2, 3, 3, 5, 5, 6, 7, 9, so the four smallest are 1, 2, 2, 3. Any ordering of those is accepted.
The statement asks for O(n log V) time, where V is the range of the values, from the smallest to the largest. Sorting would cost O(n log n), and a heap O(n log k). This lesson uses neither: it binary searches over the values themselves.
Approach: find the k-th smallest value, then collect
If you know the value of the k-th smallest number, call it kthNum, the answer is easy to assemble. Every number strictly smaller than kthNum belongs in the answer. The slots left over are filled with copies of kthNum itself.
The padding is needed because kthNum may appear several times. In the example, kthNum is 3. Only three numbers are strictly smaller (1, 2, 2), so one copy of 3 completes the four.
To find kthNum without sorting, use the function howManyLess(nums, val), which counts the numbers below val in one pass. The k-th smallest number is the largest value that has fewer than k numbers below it. In the example, 3 has three numbers below it, which is fewer than 4. The value 4 has five below it, so it is too big.
That property is monotone: as val grows, the count of smaller numbers never drops. So a binary search works, and it searches over values, not positions. The search range runs from the smallest to the largest value in nums, which is V wide, so it takes O(log V) rounds, and each round costs one O(n) pass. The middle value tried doesn't have to be in nums; the example tries 4, which isn't there.
Step by step
- Write
howManyLess. It takesnumsandvaland returns how many numbers are less thanval. - Start a counter at 0.
- Loop over
numswithrange, naming each valuenumand discarding the index. - Count the smaller ones. When
val > num, increment the counter. - Return the counter.
- Declare
kSmallestNum. It takesnumsandkand returns the k-th smallest value. - Initialize the bounds.
left, right := nums[0], nums[0]start at the first value and are about to be widened. - Scan the slice to find the extremes.
- Lower the minimum. If
num < left, setleft = num. - Raise the maximum. If
num > right, setright = num. - Keep the best answer so far.
result := leftis a safe start, since the smallest value has nothing below it, and 0 is fewer thank. - Binary search while
left <= right. This is aforwith only a condition, Go's version ofwhile. - Pick the middle value.
val := left + (right-left)/2. - Count the numbers below it.
smaller := howManyLess(nums, val). - If
smaller < k,valis still a candidate for the k-th smallest. - Save it in
result. - Look higher. Set
left = val + 1to see whether a larger value also qualifies. - Otherwise
valhaskor more numbers below it. It is too large, so setright = val - 1. - Return
resultafter the loop. - Back in
kSmallest, store the value fromkSmallestNuminkthNum. - Create the answer with
make([]int, 0, k): length 0, room forkvalues. - Loop over
numsagain. - Take the smaller ones. When
val < kthNum, appendval. - Pad if needed. Fewer than
knumbers can be strictly smaller thankthNum, so loop whilelen(result) < k... - ...appending
kthNumeach time. - Return the slice.
The solution in Go
func kSmallestNum(nums []int, k int) int {
left, right := nums[0], nums[0]
for _, num := range nums {
if num < left {
left = num
}
if num > right {
right = num
}
}
result := left
for left <= right {
val := left + (right-left)/2
smaller := howManyLess(nums, val)
if smaller < k {
result = val
left = val + 1
} else {
right = val - 1
}
}
return result
}
func howManyLess(nums []int, val int) int {
counter := 0
for _, num := range nums {
if val > num {
counter++
}
}
return counter
}
func kSmallest(nums []int, k int) []int {
kthNum := kSmallestNum(nums, k)
result := make([]int, 0, k)
for _, val := range nums {
if val < kthNum {
result = append(result, val)
}
}
for len(result) < k {
result = append(result, kthNum)
}
return result
}
The Go details worth knowing:
make([]int, 0, k). The arguments are the type, the length and the capacity. The slice starts empty but can takekelements throughappendwithout reallocating. Writingmake([]int, k)instead would createkzeros, andappendwould add after them.appendreturns the slice. It may have to move the data to a new array, so the result is always assigned back:result = append(result, val).- Multiple assignment.
left, right := nums[0], nums[0]declares two variables in one line. Readingnums[0]panics with an index out of range on an empty slice, so the function assumesnumshas at least one value. - The midpoint form.
left + (right-left)/2cannot overflow, while(left + right) / 2could if the two sums exceeded the integer range. Go'sintis 64 bits on common platforms, so it is rarely an issue here, but the form is a habit worth keeping. The division is integer division, which rounds down for positive values. - Helper order.
howManyLessis declared afterkSmallestNumin the file, which is fine: package-level functions in Go can be used before their declaration. - Result order. The values come out in their original order in
nums, then the copies ofkthNum. For the example the function returns[2 1 2 3], which is one of the accepted orderings.
Complexity
Time: O(n log V). Finding the smallest and largest value is one pass, O(n). The binary search halves a range that is V wide, so it runs about log V times, and every round calls howManyLess, which is O(n). Building the answer is two more steps: one pass over nums and at most k appends. The binary search dominates.
Space: O(k) for the answer, O(1) extra. The result slice holds k values, and the search uses a handful of integers. Pre-sizing the capacity to k avoids copying while it fills.
Edge cases and common mistakes
- Duplicates of the k-th value. For
[4, 1, 4, 4]andk = 3,kthNumis 4, only1is strictly smaller, and two copies of 4 pad the result:[1 4 4]. Using<=in the collecting pass would take all three 4s and return too many values. - All values equal.
[7, 7, 7, 7]withk = 3has nothing strictly smaller, so the padding loop builds the whole answer:[7 7 7]. k = 1. The k-th smallest is the minimum, and the answer is just that value.kequal to the length. The k-th smallest is the maximum. Every other value is below it or equal to it, and the result contains all ofnums.- Counting
<=instead of<. The condition "fewer thanknumbers below" depends on the strict comparison. Counting equal values as well would shift the answer when there are duplicates. - Forgetting to save
result. Whensmaller < kthe search moves right, so the last qualifying value must be stored before the range moves past it. - Negative numbers. The statement promises positive integers, but the method itself doesn't rely on that:
[-3, 5, -3, 0]withk = 2gives[-3 -3].
Testing
nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5],k = 4gives the values 1, 2, 2, 3 in some order.nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5],k = 1gives[1].nums = [7, 7, 7, 7],k = 3gives[7 7 7].nums = [4, 1, 4, 4],k = 3gives[1 4 4].
