Start learning free

LESSON 200 OF 338 · SORTING

Smallest K in Go: O(n log n) with sort.Ints

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.
(There are faster solutions which we will discuss in future lessons)


STEP 1 OF 2
    Ctrl + Enter to check

    GO · SORTING

    Smallest K in Go, explained

    Problem overview

    The function kSmallest receives a slice of positive integers, nums, and a count k. It must return the k smallest values, in any order.

    In the example, 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, so the four smallest are 1, 2, 2 and 3. Duplicates count separately: the value 2 appears twice in the answer because it appears twice in the input.

    The lesson asks for O(n log n) time and O(1) extra space.

    Approach: sort, then take a prefix

    Finding the smallest values one at a time, by scanning the whole slice k times, costs O(n * k). That becomes quadratic when k is close to n. Sorting does all the ordering work at once. After an ascending sort, the k smallest values are exactly the first k positions, and a slice expression picks them out without any copying.

    The problem allows any order, but a sorted prefix is a valid answer, so there is nothing more to arrange.

    Step by step

    1. Sort nums in place. sort.Ints takes a []int, reorders it ascending and returns nothing. Writing nums = sort.Ints(nums) is a compile error, because there is no value to assign.
    2. Return the prefix. The slice expression nums[:k] is the slice of the first k elements: the low bound defaults to 0, and the high bound is exclusive, so index k itself is left out.

    The solution in Go

    func kSmallest(nums []int, k int) []int {
    sort.Ints(nums)
    return nums[:k]
}

    The code uses the sort package, so the file needs import "sort". Go treats an unused import as a compile error, and a missing one as well.

    Slices share memory

    The returned slice is not a copy. It is a view onto the same backing array as nums, which is why the lesson can claim O(1) extra space. Running the function on the example gives:

    nums := []int{5, 9, 3, 6, 2, 1, 3, 2, 7, 5}
res := kSmallest(nums, 4)
fmt.Println(res)        // [1 2 2 3]
fmt.Println(nums)       // [1 2 2 3 3 5 5 6 7 9]
fmt.Println(len(res), cap(res)) // 4 10

    Two consequences follow. First, the caller's nums has been sorted, because sort.Ints works in place. Second, writing into the result writes into nums too: after res[0] = 100, nums[0] is also 100. The capacity of res is 10 here, so an append to it would overwrite nums[4] instead of allocating a new array. If a caller needs an independent result, it can copy the prefix first:

    out := make([]int, k)
copy(out, nums[:k])

    That costs O(k) extra space, so it is not what this lesson asks for, but it is the safe choice when the result outlives the input.

    Complexity

    Time: O(n log n). The sort dominates. The slice expression takes constant time, because it only builds a new slice header (pointer, length and capacity).

    Space: O(1) extra. The sort reorders the existing array, and the result shares it. The sort itself uses a small amount of stack space for its recursion, which does not count as a copy of the input.

    Edge cases and common mistakes

    • k = 0. nums[:0] is an empty slice, and printing it gives [].
    • k larger than len(nums). A slice expression is checked against the capacity, so []int{3, 1} with k = 5 panics with "slice bounds out of range [:5] with capacity 2". The lesson's inputs keep k within the length.
    • k equal to len(nums). This is valid and returns the whole sorted slice.
    • Duplicates. They are kept. For [4, 4, 1] and k = 2 the answer is [1 4].
    • Expecting the input to stay unchanged. The function reorders nums. Copy the slice before calling it if the original order matters later.

    Try these in the editor above: the example with k = 4 should print [1 2 2 3], and [4, 4, 1] with k = 2 should print [1 4].

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