LESSON 290 OF 414 · DIVIDE AND CONQUER

Smallest K in Rust: O(n) Quickselect with In-Place Partitioning

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 15

Guided

Ctrl + Enter to check

RUST · DIVIDE AND CONQUER

Smallest K in Rust, explained

Problem overview

The function k_smallest receives a vector of positive integers, nums, and a number k. It returns the k smallest values, in any order. The lesson 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 four smallest values are 1, 2, 2 and 3. Duplicates count separately: the second 2 is the third smallest. Any order is accepted, so [2, 1, 2, 3] is as correct as [1, 2, 2, 3].

Approach: partition, then follow only the side that matters

Sorting the whole vector and taking the first k elements works, but costs O(n log n), and the order of the result doesn't even matter. A faster idea borrows the partition step from Quick Sort.

Partitioning picks a pivot and rearranges a range so that everything less than or equal to it comes first, then the pivot, then the larger values. The pivot lands at its final sorted position, pivot. So the range [left, pivot] holds exactly the left_size = pivot - left + 1 smallest values of the range. Compare that with k:

left_size >= k: all k smallest values are in [left, pivot]. Search only below the pivot, for the same k.
left_size < k: every value in [left, pivot] belongs to the answer. Only k - left_size more are needed, and they come from [pivot + 1, right].

Unlike Quick Sort, only one side is processed. With a reasonable pivot each step cuts the range roughly in half, so the work is about n + n/2 + n/4 + ..., which is under 2n: that is O(n). When the process ends, the first k positions of the vector hold the k smallest values, in no particular order. Nothing past index k needs to be sorted.

Step by step

The partition

Declare partition taking nums: &mut Vec<i32> and the inclusive usize bounds left and right, returning the pivot's final index as a usize.
Choose the pivot: the last value of the range, value = nums[right]. A marker correct_index starts at left. It is the position for the next value that is not larger than the pivot.
Scan with for i in left..=right. The ..= includes right, so the pivot is visited last.
Swap small values into place. If nums[i] <= value, swap it with nums[correct_index] and move the marker forward.
Return correct_index - 1. The pivot was the last value moved, so it sits just before the marker.

The recursive helper

Declare smallest_k_helper with the vector, left, right and k. It returns nothing and rearranges nums in place so that the k smallest values of [left, right] come first in that range.
Base case. If left >= right, there is nothing to rearrange. Return.
Partition and compute left_size as pivot - left + 1.
Enough on the left. If left_size >= k, recurse on [left, pivot - 1] with the same k, but only when pivot > left. Otherwise pivot - 1 would underflow a usize at 0.
Otherwise, in an else if that checks pivot + 1 <= right, the whole left side is in the answer. Recurse on [pivot + 1, right] for the remaining k - left_size values.

The entry point

Convert k to a usize named k_usize and store the length in n.
If n > 1, call the helper on the whole vector: &mut nums, 0, n - 1 and k_usize.
Return the first k_usize elements copied into a new vector with nums[..k_usize].to_vec(), as the final expression.

The solution in Rust

impl Solution {
    fn partition(nums: &mut Vec<i32>, left: usize, right: usize) -> usize {
        let value = nums[right];
        let mut correct_index = left;
        for i in left..=right {
            if nums[i] <= value {
                nums.swap(i, correct_index);
                correct_index += 1;
            }
        }
        correct_index - 1
    }

    fn smallest_k_helper(nums: &mut Vec<i32>, left: usize, right: usize, k: usize) {
        if left >= right {
            return;
        }
        let pivot = Self::partition(nums, left, right);
        let left_size = pivot - left + 1;
        if left_size >= k {
            if pivot > left {
                Self::smallest_k_helper(nums, left, pivot - 1, k);
            }
        } else if pivot + 1 <= right {
            Self::smallest_k_helper(nums, pivot + 1, right, k - left_size);
        }
    }

    fn k_smallest(mut nums: Vec<i32>, k: i32) -> Vec<i32> {
        let k_usize = k as usize;
        let n = nums.len();
        if n > 1 {
            Self::smallest_k_helper(&mut nums, 0, n - 1, k_usize);
        }
        nums[..k_usize].to_vec()
    }
}

The Rust details worth a closer look:

mut in the signature. k_smallest(mut nums: Vec<i32>, ...) takes ownership of the vector and makes the binding mutable, so &mut nums is allowed. The caller's vector has been moved in, so the function is free to reorder it.
Slicing and copying. nums[..k_usize] is a slice of the first k_usize elements, and to_vec() copies it into an owned Vec<i32>. The line has no semicolon, so it is the return value.
Unsigned arithmetic. pivot - left + 1 is safe because pivot >= left. In k - left_size, the branch only runs when left_size < k, so the subtraction can't go below zero. Without those facts an underflow would wrap in an optimized build or panic in a debug one.
No random pivot. Always taking the last element is simple, and the standard library has no random number generator to pick one. On shuffled data the expected time is still O(n).
Passing nums along. Inside the helpers, nums is already a mutable reference, so it is passed to the next call as it is.

Complexity

Time: O(n) on average. Each level partitions its range with a single pass and then continues in only one part. With reasonably balanced splits the ranges shrink geometrically, so the passes add up to a constant times n. Because the pivot is always the last element, a vector that is already sorted, or has many equal values, splits badly: each step removes only one element, and the work becomes O(n^2).

Space: O(log n) on average. The vector is rearranged in place, and the only extra memory is the recursion. Each level of the recursion halves the range in the good case, so the depth is about log n. The result vector holds k values, which is the output.

Edge cases and common mistakes

k = 0. The result is an empty vector: nums[..0] is an empty slice.
k equal to the length. The result holds every value. For the example with k = 10, sorting the result gives [1, 2, 2, 3, 3, 5, 5, 6, 7, 9].
k larger than the length. The slice nums[..k_usize] goes past the end and the program panics. The function assumes k is at most the length of the vector.
Empty or single-element vector. The helper isn't called. [4] with k = 1 gives [4], and [] with k = 0 gives [].
Duplicates. [2, 2, 2, 2] with k = 2 gives [2, 2]. The partition uses <=, so equal values are counted correctly.
Expecting sorted output. The first k values are the right ones but not in order. For the example, the function returns [2, 1, 2, 3].
Forgetting the guards. Without pivot > left and the n > 1 check, a usize subtraction can underflow.

Testing

nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives the values 1, 2, 2 and 3 in some order.
The same vector with k = 1 gives [1].
The same vector with k = 0 gives [].
nums = [4], k = 1 gives [4].
nums = [2, 2, 2, 2], k = 2 gives [2, 2].

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

Start learning free