Start learning free

LESSON 286 OF 338 · HEAP

Smallest K Integers in Go: O(n log k) with a Max Heap

Practice O(n log k) 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 k) time and use O(k) extra space.
(There are faster solutions which we will discuss in future lessons)


STEP 1 OF 10
    Ctrl + Enter to check

    GO · HEAP

    Smallest K in Go, explained

    Problem overview

    The function kSmallest receives a slice nums of positive integers and a number k. It returns the k smallest values of nums, in any order. Duplicates count separately: if the value 2 appears twice, both copies can be in the answer.

    In the example, nums = []int{5, 9, 3, 6, 2, 1, 3, 2, 7, 5} and k = 4. Sorted, the slice is 1, 2, 2, 3, 3, 5, 5, 6, 7, 9, so the four smallest are 1, 2, 2 and 3. Any ordering of those four values is accepted, and the solution below happens to return them from largest to smallest: [3 2 2 1].

    The statement asks for O(n log k) time and O(k) extra space.

    Approach: a max heap of size k

    The easiest solution sorts the whole slice and takes the first k values. It costs O(n log n) time and does more work than needed, since the order of the other n - k values doesn't matter.

    A heap is a structure that always gives access to its extreme value, here the largest, in O(1), and inserts or removes in O(log size). Use a max heap that holds the best k candidates seen so far:

    • While the heap has fewer than k values, every new value goes in.
    • Once it is full, its top is the largest of the current candidates. A new value smaller than the top beats it, so remove the top and insert the new value. A value that isn't smaller can't be among the k smallest, and is ignored.

    The heap never holds more than k values, so each operation costs O(log k). Doing that for n values gives O(n log k), and the memory is O(k).

    On the example with k = 4, the first four values 5, 9, 3 and 6 fill the heap, and its top is 9. Then:

    • 2 is smaller than 9, so 9 leaves and 2 enters. The top is now 6.
    • 1 replaces the 6, and the top is 5.
    • 3 replaces the 5, and the top is 3.
    • The next 2 is smaller than 3, so it replaces a 3. The heap now holds 3, 2, 2 and 1, with 3 on top.
    • 7 and the last 5 are not smaller than 3, so they are ignored.

    Step by step

    1. Create the heap. Go's container/heap package works on any type that implements heap.Interface. The lesson provides a maxHeap type, a slice of ints with the greatest value on top. Create an empty one and keep a pointer to it in smallestKNums, because pushing and popping change the slice's length.
    2. Loop over nums with range, discarding the index with _ and naming the element value.
    3. While the heap is not full, that is when smallestKNums.Len() < k, push the value with heap.Push.
    4. Otherwise compare with the top. The heap is full, so look at the greatest value stored, (*smallestKNums)[0]. If the new value is smaller, pop the top with heap.Pop and push the new value. The result of the pop is not needed.
    5. Allocate the result with make([]int, k).
    6. Unload the heap. Loop k times, pop the greatest remaining value with heap.Pop, and convert it back with the type assertion .(int). The values come out from largest to smallest, which is fine because any order is accepted.
    7. Return the slice.

    The solution in Go

    The solution relies on the maxHeap type that the lesson provides. A typical definition, for reference, looks like this (you need the container/heap import to use it):

    type maxHeap []int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *maxHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

    Here is the function itself:

    func kSmallest(nums []int, k int) []int {
    smallestKNums := &maxHeap{}
    for _, value := range nums {
        if smallestKNums.Len() < k {
            heap.Push(smallestKNums, value)
        } else if value < (*smallestKNums)[0] {
            heap.Pop(smallestKNums)
            heap.Push(smallestKNums, value)
        }
    }
    result := make([]int, k)
    for i := 0; i < k; i++ {
        result[i] = heap.Pop(smallestKNums).(int)
    }
    return result
}

    The Go details worth knowing:

    • Less is reversed. Less(i, j) returns h[i] > h[j], so the "smallest" element by the heap's order is the largest number. That is what turns a min-heap implementation into a max heap.
    • Methods and package functions. The type's own Push and Pop only append and remove at the end of the slice. The package functions heap.Push and heap.Pop call them and also restore the heap order. Calling the methods directly would corrupt the heap.
    • Pointer receivers. Push and Pop change the slice's length, so they must work on a pointer. That is why smallestKNums is &maxHeap{} and why indexing needs the dereference (*smallestKNums)[0]. The parentheses are required because indexing binds tighter than *.
    • Values come back as any. heap.Pop returns an interface value, so .(int) asserts the concrete type. The assertion panics if the type is wrong, but here every element was pushed as an int.
    • Len through the pointer. smallestKNums.Len() works on a pointer because the method set of *maxHeap includes the value-receiver methods.
    • The top is index 0. The heap keeps its greatest value at the front of the slice.

    Complexity

    Time: O(n log k). Each of the n values costs at most one comparison with the top, one pop and one push, each O(log k) on a heap with at most k values. Unloading adds k pops, O(k log k), which is no more than O(n log k) when k <= n.

    Space: O(k). The heap holds at most k values, and the result has k values.

    Edge cases and common mistakes

    • Duplicates. nums = []int{4, 4, 4} and k = 2 returns [4 4]. The comparison is strict, so an equal value doesn't replace the top, and that's fine because the result is the same.
    • k equals the length. nums = []int{3, 1, 2} and k = 3 returns all the values from largest to smallest, [3 2 1].
    • k = 1. The heap holds one value, the smallest so far. nums = []int{3, 1, 2} returns [1].
    • k out of range. The solution assumes 1 <= k <= len(nums). With k = 0 and a non-empty slice, the expression (*smallestKNums)[0] reads an empty heap and panics with an index out of range error. With k greater than len(nums), the unloading loop pops from an empty heap and panics. Check these cases first if your input can contain them.
    • Using a min heap by mistake. A min heap would give the smallest value on top, which is exactly the one you want to keep. The heap needs the largest candidate on top to decide what to evict.
    • Calling the methods directly. smallestKNums.Push(value) appends without restoring the order, so the top is no longer the largest value.
    • Expecting sorted output. The values come out from largest to smallest. Sort the result if you need them ascending.

    Testing

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