Problem overview
The method kSmallest receives an array of positive integers, nums, and a number k. It returns the k smallest values, in any order. For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the result contains 1, 2, 2 and 3. Duplicates count separately, so the two 2s are both in the answer and only one of the two 3s is.
The statement asks for O(n log n) time. This lesson reaches it by binary search on the values themselves, without sorting the array.
Approach: find the k-th smallest value, then collect
If we knew the value of the k-th smallest element, call it kthNum, collecting the answer would be easy. Every element smaller than kthNum belongs to the answer. If fewer than k elements are smaller, the missing places are filled with copies of kthNum, because it may occur several times.
To find kthNum, binary search over the range of values from the minimum to the maximum of nums. For a candidate value val, count how many elements are strictly smaller than it. If that count is below k, then val is not too large to be the k-th smallest: fewer than k elements precede it. So save it as the best answer so far and try bigger values. Otherwise at least k elements are smaller than val, so it is too big, and the search goes lower. The largest value that passes the test is the k-th smallest.
In the example with k = 4, the value 3 has three smaller elements (1, 2, 2), fewer than 4, so it passes. The value 4 has five smaller elements (1, 2, 2, 3, 3), so it fails. The k-th smallest is 3.
Step by step
- Write
howManyLessabovekSmallestNum. It takes_ nums: [Int]and_ val: Intand returns anInt. It startsvar counter = 0, loops overnumswithfor num in nums, adds one whenval > num, and returnscounter. - Write
kSmallestNum, which takes_ nums: [Int]and_ k: Intand returns the value of the k-th smallest. - Find the minimum and the maximum. Start
var leftandvar rightatnums[0]and update them withminandmaxin a loop. - Start
var resultatleft. - Binary search while
left <= right: take the middle valueval, and store insmallerhow many numbers are less than it, usinghowManyLess. Both arelet, since each iteration makes new ones. - If
smaller < k, savevalinresultand setleft = val + 1. Otherwise setright = val - 1. - Return
result. - Finish
kSmallest. Put the value fromkSmallestNuminto a constantkthNum, and make an emptyvar resultof type[Int]. - Append everything smaller than
kthNum. - Pad with
kthNumwhileresult.count < k, and returnresult.
The solution in Swift
class Solution {
func kSmallestNum(_ nums: [Int], _ k: Int) -> Int {
var left = nums[0]
var right = nums[0]
for num in nums {
left = min(left, num)
right = max(right, num)
}
var result = left
while left <= right {
let val = left + (right - left) / 2
let 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 {
var counter = 0
for num in nums {
if val > num {
counter += 1
}
}
return counter
}
func kSmallest(_ nums: [Int], _ k: Int) -> [Int] {
let kthNum = kSmallestNum(nums, k)
var result = [Int]()
for val in nums {
if val < kthNum {
result.append(val)
}
}
while result.count < k {
result.append(kthNum)
}
return result
}
}
The Swift details worth a closer look:
- Two
resultvariables.kSmallestNumandkSmallesteach have a localresult, of typeIntand[Int]. They are separate and don't clash. [Int]()creates an empty array of integers, andappendadds one element at the end. The array is avarbecause it grows.- Value semantics. The
numsparameter is never modified, and passing it tohowManyLesson every iteration does not copy the elements: Swift shares the storage until someone mutates it. - Midpoint without overflow.
left + (right - left) / 2avoids adding two potentially large values. Swift traps on overflow, so for values nearInt.maxthis form matters. - An alternative count.
nums.filter { $0 < val }.countgives the same number, but builds a temporary array first.
Complexity
Let n be the length of nums and M the largest value.
Time: O(n log M). The search halves the value range [min, max] every iteration, so it runs about log M times, and each iteration counts the smaller elements in O(n). For integers that fit in a fixed width this is within the O(n log n) the statement allows. The final collection pass is O(n).
Space: O(k) for the returned array, and no other extra memory.
Edge cases and common mistakes
- Duplicates of the k-th value. For
[4, 4, 4, 4]andk = 2no element is smaller than 4, so the result is filled with two 4s. Without the padding step the result would be empty. kequals 1. The k-th smallest is the minimum, nothing is smaller, and the result is one copy of it.kequals the length of the array. The k-th smallest is the maximum, and the result holds all the elements smaller than it plus enough copies of the maximum.- A single element.
[7]withk = 1returns[7]. - Counting values that are at most
val. Using<=inhowManyLesschanges the meaning of the test and shifts the answer to the wrong value. The test must count strictly smaller elements. - Empty input.
nums[0]on an empty array crashes at run time, so the method needs at least one element.
Testing
nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5],k = 4gives the values 1, 2, 2 and 3 in some order; this code returns[2, 1, 2, 3].nums = [4, 4, 4, 4],k = 2gives[4, 4].nums = [7],k = 1gives[7].nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5],k = 1gives[1].
