Start learning free

LESSON 224 OF 415 · SORTING

Smallest K in Swift: O(n log n) with sorted() and a Slice

Practice O(n log n) Read the full explanation

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:

For this lesson, your algorithm should run in O(n log n) time and use O(1) extra space besides the returned values, when sorting in place (sorting a copy of nums takes O(n) extra space).
(There are faster solutions which we will discuss in future lessons)


STEP 1 OF 2
    Ctrl + Enter to check

    SWIFT · SORTING

    Smallest K in Swift, explained

    Problem overview

    The method kSmallest receives an array nums of positive integers and a number k. It must return the k smallest values, in any order. Duplicates count as separate values.

    For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the values in increasing order are 1, 2, 2, 3, 3, 5, 5, 6, 7, 9. The first four are 1, 2, 2 and 3, so the answer is [1, 2, 2, 3]. Any ordering of those four values is accepted.

    Approach: sort and take the front

    If the array were sorted in ascending order, the smallest k values would simply be its first k elements. So the plan has two parts: sort, then cut off the first k elements.

    Sorting costs O(n log n) time, which is the budget the lesson sets. Cutting the front is cheap next to it. Later lessons on the same problem get below O(n log n) with other tools, but sorting is the shortest correct solution.

    Step by step

    1. Sort into a new array. A Swift function parameter is a constant, so nums.sort() isn't allowed. nums.sorted() returns a sorted copy, which you store in sortedNums.
    2. Take the first k elements. Subscript the sorted array with the range 0..<k. The half-open range includes index 0 and stops before index k, which gives exactly k elements.
    3. Convert to an array. A range subscript gives an ArraySlice, not an [Int]. The function returns [Int], so the slice is wrapped in Array(...).

    The solution in Swift

    class Solution {
    func kSmallest(_ nums: [Int], _ k: Int) -> [Int] {
        let sortedNums = nums.sorted()
        return Array(sortedNums[0..<k])
    }
}

    The Swift details worth knowing:

    • sorted() returns a new array. The caller's nums is unchanged. This is also why a let is enough for sortedNums.
    • An ArraySlice is a view. sortedNums[0..<k] doesn't copy the elements. It points into sortedNums and keeps its original indices, so a slice taken from sortedNums[2..<5] starts at index 2, not 0. Wrapping it in Array(...) makes a standalone array whose indices start at 0.
    • Out-of-range slices crash. If k is larger than nums.count, the range 0..<k goes past the end, and Swift stops the program with a run-time error. The lesson expects k to be a valid count. If you wanted to be safe, Array(sortedNums.prefix(k)) returns the whole array when k is too big.
    • k = 0 is fine. The range 0..<0 is empty, so the result is [].

    Complexity

    Time: O(n log n). The sort dominates. Building the slice and the array costs O(k), which is at most O(n).

    Space: O(n). sorted() builds a second array of n elements, and the result holds k more. Because Swift arrays are values, sorting a var copy of the parameter in place would need an n-element copy too.

    Edge cases and common mistakes

    • Duplicates. Each copy of a value is a separate element. For [2, 2, 1] with k = 2 the answer is [1, 2], and for [2, 2, 1] with k = 3 it is [1, 2, 2].
    • k equal to the length. 0..<k covers the whole array, and the result is the sorted array.
    • Using 0...k. The closed range includes index k and returns k + 1 elements. It also crashes when k equals the length.
    • Calling sort() on nums. This doesn't compile, because nums is a constant. Swift reports cannot use mutating member on immutable value.
    • Returning the slice itself. A function declared to return [Int] can't return an ArraySlice<Int>. The compiler rejects it until you add Array(...).

    Testing

    • [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] with k = 4 gives [1, 2, 2, 3].
    • [4, 1] with k = 0 gives [].
    • [4, 1, 2] with k = 3 gives [1, 2, 4].
    Keep going with SwiftAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free