Start learning free

LESSON 226 OF 414 · SORTING

Smallest K Integers in Rust: O(n log n) with Sorting

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

    RUST · SORTING

    Smallest K in Rust, explained

    Problem overview

    The function k_smallest receives a vector of positive integers, nums, and a number k. It returns a vector with the k smallest values of nums. Duplicates count separately, and the result may be in any order.

    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 four smallest are 1, 2, 2 and 3, so the answer is [1, 2, 2, 3]. The second 2 is a separate element, and it counts as the third smallest.

    The lesson asks for O(n log n) time. Faster solutions exist and come later in the course, but sorting is the natural first idea.

    Approach: sort, then take a prefix

    After sorting in ascending order, the smallest value is first, the second smallest is next, and so on. The k smallest values are exactly the first k elements. No searching or counting is needed, and ties between equal values don't matter, because equal values are interchangeable.

    Sorting dominates the cost. Taking the first k elements is a copy of at most n values, which is cheaper than the sort. The total is therefore O(n log n).

    Step by step

    1. Sort nums in place with sort(). The parameter is declared mut nums, so the function owns the vector and may reorder it. The caller's vector was moved into the function, so there is no other copy to disturb.
    2. Return the first k elements as a new vector. Slice the sorted vector up to k and call to_vec(), which copies the slice into an owned Vec<i32>. Because the line has no semicolon, it is the function's return value.

    The solution in Rust

    impl Solution {
    fn k_smallest(mut nums: Vec<i32>, k: i32) -> Vec<i32> {
        nums.sort();
        nums[..k as usize].to_vec()
    }
}

    The Rust details worth a closer look:

    • Slices and ranges. nums[..k as usize] is the range from the start up to, but not including, index k, so it has k elements. It is a borrowed view of the sorted vector, not a new vector, which is why to_vec() is needed to produce something the function can return.
    • The cast. Indices and slice bounds must be usize, while k is an i32. The cast k as usize binds tighter than the range, so it converts k alone.
    • Casting is silent. For a negative k, k as usize doesn't fail: -1i32 as usize is 18446744073709551615 on a 64-bit target, and the slice then panics for being out of range. A valid k is never negative, so that case does not arise here.
    • Out-of-range slices panic. If k were larger than the length of nums, the slice would stop the program with an error such as range end index 3 out of range for slice of length 2. The lesson's inputs keep k within the length.
    • sort() returns nothing. It reorders the vector in place and gives back (), so the next line uses nums itself, not a returned value.
    • sort_unstable() would work as well. For integers, equal values can't be told apart, so not keeping their original order costs nothing, and it is often a little faster.

    Complexity

    Time: O(n log n). Sorting n integers takes O(n log n), and copying k of them takes O(k) with k <= n.

    Space. The vector is sorted in place, and the returned vector holds k elements. Rust's stable sort() may allocate a temporary buffer, so strictly its extra space can be O(n), while sort_unstable() sorts without one.

    Edge cases and common mistakes

    • k equal to 0. The slice nums[..0] is empty, so the function returns [].
    • k equal to the length. The slice covers the whole vector, so the answer is the entire sorted vector. For [5, 9] and k = 2 the result is [5, 9].
    • Duplicates. With [2, 2, 2, 1] and k = 3 the answer is [1, 2, 2]. Each copy is its own element, so values are not removed when they repeat.
    • Slicing before sorting. Taking the first k elements of the unsorted vector gives the first k values in the input, not the smallest ones.
    • Sorting in the wrong direction. Sorting in descending order and then taking a prefix would give the largest k values.
    • Using k as an index. The slice nums[..k] without the cast doesn't compile, because k is an i32 and a slice bound must be usize.

    Testing

    • nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives [1, 2, 2, 3].
    • nums = [5, 9], k = 0 gives [].
    • nums = [5, 9], k = 2 gives [5, 9].
    • nums = [2, 2, 2, 1], k = 3 gives [1, 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