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. Duplicates count separately.
In the example, nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4. Sorted, the array begins 1, 2, 2, 3, so the answer is those four values. The order doesn't matter, so [2, 1, 3, 2] is also correct.
The lesson asks for O(n log k) time and O(k) extra space.
Approach: keep the k smallest seen so far
Sorting the whole array and taking the first k values works, but it costs O(n log n) time and ignores that k may be small. A heap does better by remembering only k values at any moment.
The trick is to use a max heap, where the greatest element sits on top. Read the numbers one by one:
- While the heap holds fewer than
kvalues, add the new one. - Once it is full, the top is the greatest of the
ksmallest values so far. A new value smaller than the top belongs in the set, and the top is the one that no longer does. Remove the top and add the new value. - A new value that is not smaller than the top changes nothing.
Each push or pop costs O(log k) because the heap never holds more than k values, and there are at most two per element, so the loop costs O(n log k). At the end the heap holds exactly the answer.
Swift's standard library has no heap type. The lesson supplies a MaxHeap struct with push, pop, peek and count.
Step by step
- Create the heap.
var smallestKNums = MaxHeap(). It is avarbecausepushandpopchange the struct. - Loop over the values.
for value in nums. - Fill the heap first. If
smallestKNums.count < k, pushvalue. - Compare with the top when full. Otherwise write
} else if value < smallestKNums.peek(). The methodpeek()reads the top without removing it. - Swap the values. Inside that branch, call
pop()to remove the greatest value, assigning its result to_because it isn't needed, then pushvalue. The heap again holds exactlykvalues. - Empty the heap into an array. Declare
var result = [Int](). WhilesmallestKNums.count > 0, append the result ofpop()to it. - Return
result. The values come out from largest to smallest, which is fine because any order is accepted.
The solution in Swift
struct MaxHeap {
private var data = [Int]()
var count: Int { return data.count }
func peek() -> Int { return data[0] }
mutating func push(_ value: Int) {
data.append(value)
var i = data.count - 1
while i > 0 {
let parent = (i - 1) / 2
if data[parent] >= data[i] { break }
data.swapAt(parent, i)
i = parent
}
}
mutating func pop() -> Int {
let top = data[0]
let last = data.removeLast()
if !data.isEmpty {
data[0] = last
var i = 0
let size = data.count
while true {
let left = 2 * i + 1
let right = 2 * i + 2
var largest = i
if left < size && data[left] > data[largest] { largest = left }
if right < size && data[right] > data[largest] { largest = right }
if largest == i { break }
data.swapAt(i, largest)
i = largest
}
}
return top
}
}
class Solution {
func kSmallest(_ nums: [Int], _ k: Int) -> [Int] {
var smallestKNums = MaxHeap()
for value in nums {
if smallestKNums.count < k {
smallestKNums.push(value)
} else if value < smallestKNums.peek() {
_ = smallestKNums.pop()
smallestKNums.push(value)
}
}
var result = [Int]()
while smallestKNums.count > 0 {
result.append(smallestKNums.pop())
}
return result
}
}
The Swift details worth a closer look:
- The heap is an array. The element at index
ihas its children at2 * i + 1and2 * i + 2and its parent at(i - 1) / 2.pushadds at the end and moves the value up while it is larger than its parent.popmoves the last element to the root and sinks it down, swapping with the larger child. mutating. A struct is a value type, so a method that changesdatamust be markedmutating, and it can only be called on avar.peek()andcountonly read, so they need neither.private var data. Code outside the struct can't touch the array, so onlypushandpopcan break the heap order._ = smallestKNums.pop().pop()returns a value and isn't marked to allow ignoring it, so a bare call produces a compiler warning about an unused result. Assigning to_says the discard is deliberate.- Value semantics. Assigning a
MaxHeapto another variable copies it, so the two heaps then change independently. swapAt.data.swapAt(i, j)exchanges two elements in place.
Complexity
Time: O(n log k). Each of the n values costs a comparison with the top and, at most, one pop and one push on a heap of size k, each O(log k). Draining the heap at the end costs k pops, which is O(k log k) and no more than the main loop.
Space: O(k). The heap never holds more than k values, and the result array also has k.
Edge cases and common mistakes
- Duplicates. The comparison is strict, so a value equal to the top is skipped, which is fine because swapping equal values changes nothing. For
[4, 4, 4]withk = 2the result is[4, 4]. klarger than the array. The heap never fills, every value is pushed, and the method returns all values.[5, 9, 3]withk = 5gives[9, 5, 3].- Empty array. The method returns an empty array.
k = 0. The reference does not handle it: the heap stays empty, and the first value reachespeek(), which readsdata[0]of an empty array and crashes the program. Guard against it withif k <= 0 { return [] }whenkcan be 0.- Using a min heap. The top would be the smallest value, and you would be discarding the wrong element. The set of
ksmallest values needs its greatest member on top. - Expecting sorted output. The values come out in descending order. Sort the result if you need ascending order.
Testing
nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5],k = 4gives[3, 2, 2, 1], which has the same values as the example.nums = [5, 9, 3],k = 5gives[9, 5, 3].nums = [4, 4, 4],k = 2gives[4, 4].nums = [],k = 3gives[].
