How do I repeat some elements in a vector based on a condition?

Question

I encountered this problem during a kata. My more readable implementation was the following:

use std::vec::Vec;

fn repeat_even(v: Vec<i32>) -> Vec<i32> {
    v.into_iter().flat_map(|x| match x % 2 { 0 => vec![x, x], _ => vec![x] }).collect()
}

fn main() {
    let v = vec![1, 2, 3, 4, 6];
    assert_eq!(repeat_even(v), vec![1, 2, 2, 3, 4, 4, 6, 6]);
}

I have two questions about it:

• Is it necessary to create another Vec? Is is possible to use the same Vec, i.e. modify it while iterating?

• My solution is, I imagine, inefficient: I allocate a lot of vectors, and I have no guarantee that this will be optimized. Is there a better solution: readable and with less allocation?

Answer 1

flat_map expects iterators, so you can return an iterator of the values:

use std::iter;

fn double_even(v: &[i32]) -> Vec<i32> {
    v.iter().flat_map(|&x| {
        let count = if x % 2 == 0 { 2 } else { 1 };
        iter::repeat(x).take(count)
    }).collect()
}

fn main() {
    let v = vec![1, 2, 3, 4, 6];
    assert_eq!(double_even(&v), vec![1, 2, 2, 3, 4, 4, 6, 6]);
}

Things to note:

• There's no reason to have use std::vec::Vec. It's already imported via the prelude.
• You aren't making use of the memory allocation of the passed-in vector, so there's no reason to take it. See also Why is it discouraged to accept a reference to a String (&String) or Vec (&Vec) as a function argument?
• Don't use vec!(); use vec![] instead. It doesn't matter to the compiler, but it matters to humans.

---

If you were really set on attempting to reuse the memory, I'd walk backwards along the iterator to help avoid index invalidation:

fn double_even(mut v: Vec<i32>) -> Vec<i32> {
    for i in (0..v.len()).rev() {
        let val = v[i]; 
        if val % 2 == 0 {
            v.insert(i, val);
        }
    }
    v
}

This is probably algorithmically worse; each insert moves all the data after it. I believe the worst-case would be O(n^2) when every element were even.

I also wouldn't normally take by-value here. I'd instead take a mutable reference. You could always wrap it back in a value if you really needed it:

fn double_even_ref(v: &mut Vec<i32>) {
    for i in (0..v.len()).rev() {
        let val = v[i];
        if val % 2 == 0 {
            v.insert(i, val);
        }
    }
}

fn double_even(mut v: Vec<i32>) -> Vec<i32> {
    double_even_ref(&mut v);
    v
}