LESSON 317 OF 414 · HEAP

Smallest K Integers in Rust: 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 7

Guided

Ctrl + Enter to check

RUST · HEAP

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. Duplicates count separately.

In the example, nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4. Sorted, the vector is [1, 2, 2, 3, 3, 5, 5, 6, 7, 9], so the four smallest are 1, 2, 2 and 3. The result may come in any order, so [2, 1, 3, 2] is just as correct.

The statement sets the budget: O(n log k) time and O(k) extra space.

Approach: keep the k best so far, with the worst on top

The obvious way is to sort the whole vector and take the first k elements. That costs O(n log n) time and needs a copy of the data or a mutable input. The budget asks for something that depends on k, which can be much smaller than n.

The idea is to keep a collection of the k smallest values seen so far, and to be able to find the largest of them quickly. That largest value is the one to throw out when a smaller one arrives. A max-heap does exactly this: it gives the greatest element in O(1) and inserts or removes in O(log k) when it holds at most k elements.

The loop looks at each value once:

While the heap has fewer than k elements, push the value.
Once it is full, compare the value with the heap's top, the largest of the kept values. If the new value is smaller, it belongs among the k smallest and the top does not. Remove the top and push the new value.
If the new value is not smaller, skip it.

When the loop ends, the heap holds exactly the k smallest values. Each of the n values costs at most one pop and one push, each O(log k), which gives O(n log k).

Rust's BinaryHeap is a max-heap by default, which is the one needed here, so no custom ordering is required.

Step by step

Convert k to usize by shadowing it with a new let k. It will be compared with heap.len(), which is a usize, and Rust doesn't compare different integer types.
Create the heap as a mutable BinaryHeap<i32> with BinaryHeap::new().
Loop over &nums with the pattern &value. Iterating over a reference to a vector yields &i32 items, and the pattern copies each number out, so value is a plain i32.
Fill the heap first. If heap.len() < k, push the value.
When the heap is full, read its top. heap.peek() returns an Option holding a reference to the greatest element, or None for an empty heap. The pattern Some(&top) unwraps it and copies the number into top.
Replace the top when the new value is smaller. If value < top, call heap.pop() to discard the largest, then push the new value. The heap keeps exactly k elements.
Return the heap's contents. heap.into_iter().collect(), as the final expression with no semicolon, builds a Vec<i32>, because that is the function's return type.

The solution in Rust

use std::collections::BinaryHeap;

impl Solution {
    fn k_smallest(nums: Vec<i32>, k: i32) -> Vec<i32> {
        let k = k as usize;
        let mut heap: BinaryHeap<i32> = BinaryHeap::new();
        for &value in &nums {
            if heap.len() < k {
                heap.push(value);
            } else if let Some(&top) = heap.peek() {
                if value < top {
                    heap.pop();
                    heap.push(value);
                }
            }
        }
        heap.into_iter().collect()
    }
}

The Rust details worth a closer look:

Why copy the top out. peek() borrows the heap. If top stayed a reference, the heap would remain borrowed, and heap.pop() inside the if would not compile. Writing Some(&top) copies the number, which ends the borrow. A line such as println!("{:?} {:?}", h.peek(), h.pop()) fails for the same reason: the compiler says h cannot be borrowed as mutable because it is also borrowed as immutable.
Ignoring the result of pop(). It returns an Option with the removed value. The call is a statement, so the value is dropped. That is what is wanted: the largest element is simply discarded.
The order of the result. into_iter() on a BinaryHeap yields the elements in an unspecified order, not sorted. The statement accepts any order. To get them sorted, use into_sorted_vec(), which returns ascending order.
as casts. k as usize doesn't check the value. A negative k would wrap to a huge number: -1i32 as usize is 18446744073709551615, and the heap would keep every element. The lesson's inputs don't have negative values for k.
Shadowing. let k = k as usize; creates a new variable with the same name and a different type. The old k can't be used after it, and that is the point.
else if let. It combines an else with a pattern test. The block runs only when the pattern matches. Here the heap is full, so with k > 0 the pattern always matches. With k = 0 it is the pattern that fails on an empty heap, which is why nothing is ever pushed.

Complexity

Let n be the number of values.

Time: O(n log k). The loop runs n times. In each iteration the work is a length check, a peek and a comparison, all O(1), plus at most one pop and one push on a heap of at most k elements, each O(log k). Turning the heap into a vector at the end is O(k).

Space: O(k). The heap never holds more than k elements, and the returned vector has k elements too. The input vector isn't copied.

Edge cases and common mistakes

k is 0. Nothing is pushed and the function returns an empty vector.
k is at least nums.len(). The heap never fills up, and all values come back. With k = 15 and ten values, the result holds all ten.
Duplicates. They count separately. In the example, the result contains 2 twice, because there are two 2s among the four smallest.
Using the wrong kind of heap. A min-heap would hold the smallest values on top and make it hard to find the one to discard. The max-heap keeps the largest of the kept values on top.
Comparing with <=. Replacing the top when value == top gives the same values but does extra work for nothing. Strictly smaller is the right test.
Expecting sorted output. Don't compare against [1, 2, 2, 3] directly. Sort the result first if the order matters to you.

Testing

Sort each result before reading it, since the order is unspecified.

nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives [1, 2, 2, 3].
The same nums with k = 1 gives [1].
The same nums with k = 0 gives [].
The same nums with k = 15 gives all ten values: [1, 2, 2, 3, 3, 5, 5, 6, 7, 9].
nums = [], k = 3 gives [].

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

Start learning free