Rust Interview Puzzles: Sort a binary vector in linear time

Rust Interview Puzzles: Sort a binary vector in linear time

Problem

We need to sort a provided binary vector vec (a vector that contains only 0 or 1 numbers) in linear time.

Example 1:

Input: vec = [1, 0, 1, 0, 1, 0, 0, 1]
Output: [0, 0, 0, 0, 1, 1, 1, 1]

Example 2:

Input: vec = [0, 0, 1, 0, 1, 1, 0, 1, 0, 0]
Output: [0, 0, 0, 0, 0, 0, 1, 1, 1, 1]

Example 3:

Input: vec = [0, 0, 1, 1]
Output: [0, 0, 1, 1]

Solution

fn sort_binary_vector_in_linear_time(vec: &mut Vec<u8>) -> &Vec<u8> {
    let count_0 = vec.iter().fold(0, |acc, e| match e {
        0 => acc + 1,
        1 => acc,
        _ => panic!("The vector is not binary"),
    });

    for i in 0..vec.len() {
        if i < count_0 {
            vec[i] = 0;
        } else {
            vec[i] = 1;
        }
    }

    return vec;
}

Explanation

Since we know that the vector vec is a binary vector as per the requirement (could contain only values 0 and 1), we could avoid using more general sorting algorithms. Instead, we could just count a number of 0 elements (or alternatively a number of 1) using a simple fold operation. And then we iterate once again over the vector vec to fill elements with the index below that count_0 with 0, and the rest of the elements with 1.

Complexity

The time complexity of this solution is O(n) (n is a size of the input vector).


Rust Playground with this code