LESSON 257 OF 338 · DIVIDE AND CONQUER

Smallest K in Go: Quickselect in O(n) Expected Time

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 and use O(log n) extra space.

STEP 1 OF 14

Guided

Ctrl + Enter to check

GO · DIVIDE AND CONQUER

Smallest K in Go, explained

Problem overview

The function kSmallest(nums, k) receives a slice of positive integers and a number k, and 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 answer contains 1, 2, 2 and 3. A result such as [2, 1, 3, 2] is just as valid as [1, 2, 2, 3].

Approach: partition, then keep only the side that matters

Sorting the whole slice and taking the first k values costs O(n log n). The problem allows any order, which means a full sort does more work than needed.

The idea comes from Quick Sort. A partition picks a pivot and rearranges the range so that smaller-or-equal values come before the pivot and larger values after it, and it tells us the pivot's final index. If the pivot's index is pivot, the range from left to pivot holds exactly the pivot - left + 1 smallest values of the range. Compare that count with k:

If the count is at least k, all k smallest values lie on the left. Continue there, with the same k.
If the count is less than k, then everything up to the pivot belongs to the answer. Continue on the right and look for the remaining k - (pivot - left + 1) values.

Unlike Quick Sort, only one side is processed. With random pivots the range shrinks geometrically on average, so the work is n + n/2 + n/4 + ..., which adds up to O(n) expected time. This technique is called quickselect.

Take the example with k = 4. Suppose the first pivot is the last element, 5. The partition leaves [5, 3, 2, 1, 3, 2, 5, 6, 7, 9] and returns index 6. The count is 7, which is at least 4, so the search continues in indices 0 to 5. Every later step works inside that range, and when the recursion ends the first four positions hold some arrangement of 1, 2, 2 and 3.

Step by step

Declare smallestKNums above kSmallest. It takes nums []int and left, right, k int, returns nothing, and rearranges nums in place so the k smallest values of the range come first.
Base case. If left >= right, there is at most one value and nothing to do.
Declare partitionForSmallest with nums, left and right, returning the pivot's final index.
Choose a random pivot with rand.Intn(right - left + 1) added to left. Random pivots keep the expected time linear, even on sorted input.
Move it to the end with a tuple assignment.
Set up value (the pivot) and correctIndex (starting at left) in one := statement.
Scan from left to right, including right.
Move small values to the front. If nums[i] <= value, swap it with nums[correctIndex] and increase correctIndex. In Go, ++ is a statement, so it goes on its own line.
Return correctIndex - 1, where the pivot ended up.
Use the partition. Back in smallestKNums, store the returned index in pivot.
If the count pivot - left + 1 is at least k, recurse on left to pivot - 1 with the same k.
Otherwise recurse on pivot + 1 to right with k - (pivot - left + 1).
Call it from kSmallest over the whole slice, from 0 to len(nums) - 1.
Return the first k elements as nums[:k].

The solution in Go

func smallestKNums(nums []int, left, right, k int) {
	if left >= right {
		return
	}
	pivot := partitionForSmallest(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 partitionForSmallest(nums []int, left, right int) int {
	randomIndex := left + rand.Intn(right-left+1)
	nums[randomIndex], nums[right] = nums[right], nums[randomIndex]
	value, correctIndex := nums[right], left
	for i := left; i <= right; i++ {
		if nums[i] <= value {
			nums[i], nums[correctIndex] = nums[correctIndex], nums[i]
			correctIndex++
		}
	}
	return correctIndex - 1
}

func kSmallest(nums []int, k int) []int {
	smallestKNums(nums, 0, len(nums)-1, k)
	return nums[:k]
}

Go details worth knowing:

Slicing doesn't copy. nums[:k] is a new slice header over the same array, with length k and the capacity of nums. The returned slice shares memory with the input. Changing its first element also changes nums, and the caller's slice has already been rearranged by the partitions.
The high bound is exclusive. nums[:k] holds indices 0 to k-1, which is exactly k values.
Slice bounds are checked. If k is larger than len(nums), nums[:k] panics with a slice-bounds error. The problem assumes a valid k.
The import. rand.Intn belongs to math/rand, which a complete program must import.
Multiple assignment. value, correctIndex := nums[right], left declares two variables at once, and a, b = b, a swaps two elements without a temporary.
Helpers in any order. The three functions can be declared in any order in the file.

Complexity

Time: expected O(n). Each call does one pass over its range and then continues on a single side. With random pivots the range shrinks by a constant factor on average, so the passes add up to a geometric series of about 2n steps. The worst case is O(n²), which random pivots make very unlikely.

Space: expected O(log n) for the recursion. No extra slices are created, because the partition rearranges the input.

Edge cases and common mistakes

k equal to len(nums). The result is the whole slice, in some order. [3, 1, 2] with k = 3 returns the values 1, 2 and 3.
k = 1. The function returns a slice holding the minimum.
Duplicates. With [4, 4, 4] and k = 2 the answer is [4, 4]. Because values equal to the pivot go to its left, the recursion still finishes, but slices with many equal values make the partitions uneven and slow the search down.
Not subtracting the left part. When recursing right, k must shrink by the number of values already accepted. Passing the old k would look for too many.
Expecting sorted output. Only the first k positions are guaranteed to hold the right values, in no particular order. Sort the result if the order matters.
Input modified. The original slice is rearranged. Copy it first if the caller needs the order preserved.

Testing

For [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] with k = 4, the result holds 1, 2, 2 and 3 in some order. For [3, 1, 2] with k = 1 it holds 1. For [4, 4, 4] with k = 2 it is [4, 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