LESSON 282 OF 414 · BINARY SEARCH

Smallest K in Rust: Binary Search on Values

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 V) time, where V is the range of values in nums (from the smallest to the largest value).

STEP 1 OF 15

Guided

Ctrl + Enter to check

RUST · BINARY SEARCH

Smallest K in Rust, explained

Problem overview

The function k_smallest takes a vector nums of positive integers and a number k, and returns the k smallest values, in any order. For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4 the result is [1, 2, 2, 3]. Duplicates count separately: the two 2s are the 2nd and 3rd smallest values.

Any order is accepted, so [2, 1, 3, 2] is also a correct answer. The statement allows O(n log n) time for this lesson.

Approach: find the k-th smallest value, then collect

This lesson doesn't sort and doesn't use a heap. It uses binary search on the values themselves, in two parts.

First, find the value of the k-th smallest element, kth_num. For any candidate value val, a single pass can count how many elements are strictly smaller than it. If fewer than k elements are smaller than val, then val is at most the k-th smallest value, so it is a valid candidate, and the true answer could be larger, so search higher. If k or more elements are smaller, val is too big. The largest candidate that passes is the k-th smallest value.

Second, build the result. Every element strictly smaller than kth_num is certainly among the k smallest. There may be fewer than k of them, because kth_num can occur several times. The missing slots are then filled with copies of kth_num.

In the example, the counts of elements below each value are: below 1 there are 0, below 2 there is 1, below 3 there are 3, below 4 there are 5. For k = 4 the candidates 1, 2 and 3 pass (counts 0, 1 and 3 are all below 4) and 4 fails (count 5), so kth_num is 3. Elements smaller than 3 are 2, 1, 2, so one 3 is added: [2, 1, 2, 3].

Step by step

Write how_many_less(nums, val) taking a slice and a value and returning an i32: filter the elements below val, count them, and convert the usize count with as i32.
Write k_smallest_num(nums, k), the binary search on values. left and right start at the smallest and largest elements, and result starts at left.
Loop while left <= right. Take the middle value val and count the elements below it with Self::how_many_less.
If that count is below k, save val in result and search above it with left = val + 1. Otherwise search below with right = val - 1.
Return result, the value of the k-th smallest element.
In k_smallest, store that value in kth_num.
Collect the elements below kth_num into a mutable Vec<i32> named result.
Pad with kth_num in a while loop until result holds exactly k elements, then return it.

The solution in Rust

impl Solution {
    fn how_many_less(nums: &[i32], val: i32) -> i32 {
        nums.iter().filter(|&&x| x < val).count() as i32
    }

    fn k_smallest_num(nums: &[i32], k: i32) -> i32 {
        let mut left = *nums.iter().min().unwrap();
        let mut right = *nums.iter().max().unwrap();
        let mut result = left;
        while left <= right {
            let val = left + (right - left) / 2;
            let smaller = Self::how_many_less(nums, val);
            if smaller < k {
                result = val;
                left = val + 1;
            } else {
                right = val - 1;
            }
        }
        result
    }

    fn k_smallest(nums: Vec<i32>, k: i32) -> Vec<i32> {
        let kth_num = Self::k_smallest_num(&nums, k);
        let mut result: Vec<i32> = nums.iter().filter(|&&v| v < kth_num).copied().collect();
        while result.len() < k as usize {
            result.push(kth_num);
        }
        result
    }
}

Rust details worth a look:

The double reference pattern. iter() yields &i32, and filter passes its closure a reference to that item, a &&i32. The pattern |&&x| strips both layers so x is a plain i32.
copied() before collect(). After filter the items are still &i32. copied() turns them into i32 values so they collect into a Vec<i32>.
count() returns a usize, so the helper converts it with as i32. For the same reason, the length is compared with k as usize in the padding loop.
min() and max() return an Option<&i32>. unwrap() extracts the reference (and panics on an empty vector), and * copies the number.
Calling a sibling function. k_smallest_num and how_many_less have no self, so they are called with Self::. Passing &nums, a reference to a Vec, works where a slice is expected.

Complexity

Time. Each probe of the binary search counts the elements below a value in O(n). The search range runs from the minimum to the maximum value, so there are about log₂(max - min) probes. That makes O(n * log(max - min)) for the search, and the final collection is another O(n). For 32-bit values the number of probes is at most 32, so this stays within the lesson's O(n log n) budget.

Space: O(k) beyond the input, for the returned vector. The search itself uses a few integers.

Edge cases and common mistakes

Duplicates of the k-th value. With [4, 4, 4, 4] and k = 2 nothing is strictly smaller than 4, so the padding loop supplies both elements: [4, 4]. Without padding the result would be empty.
k = 1. The k-th value is the minimum: on the example it returns [1].
k = 0. No candidate passes, result stays at the minimum, nothing is strictly smaller, and the padding loop does nothing, so the result is empty.
k equal to the length. The k-th value is the maximum, and the function returns every element, in the original order apart from the copies of the maximum, which move to the end.
Changing the strict comparison. The padding step assumes that everything collected is strictly below kth_num and that the missing slots are copies of it. Counting with <= in one place and < in another breaks that match.
Returning only the elements below kth_num. That loses the copies of kth_num itself.

Testing

[5, 9, 3, 6, 2, 1, 3, 2, 7, 5] with k = 4 gives [2, 1, 2, 3], which is [1, 2, 2, 3] in a different order.
[4, 4, 4, 4] with k = 2 gives [4, 4].
[7] with k = 1 gives [7].
The example with k = 1 gives [1].

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

Start learning free