LESSON 288 OF 415 · DIVIDE AND CONQUER

Smallest K Integers in Swift: O(n) Expected Time with Quickselect

Given an array of positive integers nums, return the smallest k values, in any order you want.

Example:

Input: nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4
Output: [1, 2, 2, 3]
Explanation: Smallest number is 1, 2nd smallest is 2, 
            3rd smallest is 2, 4th smallest is 3

The result can be in any order, [2, 1, 3, 2] is also a correct answer.

Note:

Your algorithm should run in O(n) time on average and use O(log n) extra space on average (for the recursion), when partitioning in place.

STEP 1 OF 18

Guided

Ctrl + Enter to check

SWIFT · DIVIDE AND CONQUER

Smallest K in Swift, explained

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. The statement asks for O(n) time and O(log n) extra space.

For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the sorted array starts 1, 2, 2, 3, so those four values are the answer. Because any order is accepted, [2, 1, 3, 2] is correct as well. Duplicates count separately: the 2 appears twice.

Approach: partition only the side that matters

Sorting the whole array and taking the first k values costs O(n log n). This lesson doesn't need a full sort, because the answer is a set, not an ordered list.

The tool is the same partition step Quick Sort uses. Pick a pivot, and rearrange the range so that values not larger than the pivot come first, then the pivot, then the larger ones. If the pivot lands at index pivot in the range starting at left, then exactly pivot - left + 1 values (the pivot and everything before it) are the smallest of the range. Compare that count with k:

If the count is at least k, the k smallest values are all inside [left, pivot]. Continue on [left, pivot - 1] with the same k. The pivot is the largest value in its group, so if the count equals k, there is nothing more to do.
Otherwise, the whole of [left, pivot] is part of the answer, and the search continues on [pivot + 1, right] for the remaining k - (pivot - left + 1) values.

Quick Sort recurses into both sides. This method recurses into one, and with a random pivot each step shrinks the range geometrically on average. The work is about n + n/2 + n/4 + ..., which is O(n) expected. This technique is called quickselect.

An example: with [4, 1, 3] and k = 2, suppose the pivot value is 3. The values 1 and 3 are not larger than it, so the pivot lands at index 1 and the count is 2, which is at least k. The method recurses into [0, 0], which holds a single value, and stops. The first two positions hold 1 and 3.

Step by step

Declare smallestKNums with an inout [Int], the bounds left and right, and k. It rearranges the array so the k smallest values of the range come first, and returns nothing.
Stop on tiny ranges: when left >= right, there is at most one value, so return.
Declare partition with the array and bounds. It returns the pivot's final index.
Pick a random pivot: left + Int.random(in: 0...(right - left)). A random pivot keeps the expected time linear even on sorted input.
Write swapVals, which swaps two elements through a temporary constant. Its name avoids confusion with Swift's global swap.
Move the pivot to right with swapVals.
Set up the scan: value is the pivot value (a let) and correctIndex starts at left (a var).
Loop with var i from left through right in a while loop, adding 1 to i at the end of each pass.
Move small values left: if nums[i] <= value, swap it with nums[correctIndex] and add 1 to correctIndex.
Return correctIndex - 1. The pivot was the last value moved, so it sits just before correctIndex.
Partition the range in smallestKNums and save the result as pivot.
Branch on the count. If pivot - left + 1 >= k, recurse on [left, pivot - 1] with the same k. Otherwise recurse on [pivot + 1, right] with k - (pivot - left + 1).
In kSmallest, copy the parameter into var arr, run smallestKNums on the copy, and return the first k elements as an array.

The solution in Swift

class Solution {
    func swapVals(_ nums: inout [Int], _ i: Int, _ j: Int) {
        let temp = nums[i]
        nums[i] = nums[j]
        nums[j] = temp
    }
    func partition(_ nums: inout [Int], _ left: Int, _ right: Int) -> Int {
        let randomIndex = left + Int.random(in: 0...(right - left))
        swapVals(&nums, randomIndex, right)
        let value = nums[right]
        var correctIndex = left
        var i = left
        while i <= right {
            if nums[i] <= value {
                swapVals(&nums, i, correctIndex)
                correctIndex += 1
            }
            i += 1
        }
        return correctIndex - 1
    }
    func smallestKNums(_ nums: inout [Int], _ left: Int, _ right: Int, _ k: Int) {
        if left >= right {
            return
        }
        let pivot = partition(&nums, left, right)
        if pivot - left + 1 >= k {
            smallestKNums(&nums, left, pivot - 1, k)
        } else {
            smallestKNums(&nums, pivot + 1, right, k - (pivot - left + 1))
        }
    }
    func kSmallest(_ nums: [Int], _ k: Int) -> [Int] {
        var arr = nums
        smallestKNums(&arr, 0, arr.count - 1, k)
        return Array(arr[0..<k])
    }
}

The Swift details worth a closer look:

Parameters are constants. kSmallest can't sort its own nums parameter, and trying to write into it gives cannot assign through subscript: 'nums' is a 'let' constant. Copying into var arr solves it. Arrays are value types, so the copy is independent: the caller's array is not changed. Swift copies the storage lazily, when arr is first modified.
inout and &. The helper methods take the array as inout so they can rearrange it, and each call passes it as &arr or &nums.
Slices versus arrays. arr[0..<k] is an ArraySlice<Int>, a view into arr. The function must return [Int], so Array(...) converts the slice into a real array.
Half-open ranges. 0..<k excludes k, which gives exactly k elements, while 0...(right - left) includes its upper end, so the random offset can reach right - left.
Why a while loop. The loop could be written as for i in left...right. The lesson uses while with a manual i += 1, which does the same thing here.

Complexity

Time: O(n) expected. The first partition costs n. After it, the range holds fewer values, on average about half, and each later partition works on a range that is still smaller. The costs add up to a constant multiple of n. A very unlucky run of pivots could take O(n²), but random pivots make that rare.

Space: O(log n) expected. The array is rearranged in place, apart from the single copy that kSmallest makes of its input. The recursion depth follows the number of partitions, about log n on average.

Edge cases and common mistakes

k equal to the length. Every value is in the answer. [4, 1, 3] with k = 3 returns all three, in some order.
k = 0. The branch pivot - left + 1 >= k is always true, so the search keeps shrinking to the left and ends. Array(arr[0..<0]) is an empty array.
k larger than the length. The slice arr[0..<k] goes past the end, and Swift stops the program with an index-out-of-range error. The statement implies k is within the array.
Duplicates. [4, 4, 4] with k = 2 returns [4, 4]. Equal values go to the left side of the pivot because of <=, which keeps the counts exact.
Single element. [1] with k = 1 returns at once because left >= right, and the result is [1].
Expecting sorted output. Only the first k positions are guaranteed to hold the right values. Their order is not defined, so compare as sets, or sort the result first.
Recursing on both sides. That would turn the method into Quick Sort and cost O(n log n).

Testing

[5, 9, 3, 6, 2, 1, 3, 2, 7, 5] with k = 4 gives [1, 2, 2, 3] after sorting the result.
[4, 4, 4] with k = 2 gives [4, 4].
[4, 1, 3] with k = 0 gives [].
[4, 1, 3] with k = 3 gives [1, 3, 4] after sorting.
[1] with k = 1 gives [1].

Keep going with SwiftAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free