CODEWITHSUNDEEP
rustvector
Timmmm
Timmmm

Reputation: 96586

How can I retain vector elements with their original index?

If I have a Vec I can iterate over elements using an index via v.iter().enumerate(), and I can remove elements via v.retain(). Is there a way to do both at once?

In this case the index could no longer be used to access the element - it would be the index of the element before the loop was started.

I can implement this myself but to be as efficient as .retain() I'd need to use unsafe, which I'd like to avoid.

This is the result I want:

let mut v: Vec<i32> = vec![1, 2, 3, 4, 5, 4, 7, 8];

v.iter()
    .retain_with_index(|(index, item)| (index % 2 == 0) || item == 4);

assert(v == vec![1, 3, 4, 5, 4, 7]);

Upvotes: 5

Views: 3357

Answers (3)

Hauleth
Hauleth

Reputation: 23566

If you want to enumerate, filter (retain), and then collect the resulting vector then I would say to do exactly that:

v.iter()
    .enumerate()
    .filter(|&(idx, &val)| val - idx > 0)
    .collect()

Upvotes: 0

Timmmm
Timmmm

Reputation: 96586

I found essentially the same question on the Rust User's Forum. They suggested this solution, which isn't too bad:

let mut index = 0;
v.retain(|item| {
    index += 1;
    ((index - 1) % 2 == 0) || item == 4
});

At the time it wasn't a valid solution because the iteration order of retain() was not guaranteed, but happily for me someone in that thread documented the order so now it is. :-)

Upvotes: 4

richardpringle
richardpringle

Reputation: 791

@Timmmm's and @Hauleth's answers are quite pragmatic I wanted to provide a couple of alternatives.

Here's a playground with some benchmarks and tests: https://play.rust-lang.org/?version=nightly&mode=debug&edition=2018&gist=cffc3c39c4b33d981a1a034f3a092e7b

This is ugly, but if you really want a v.retain_with_index() method, you could do a little copy-pasting of the retain method with a new trait:

trait IndexedRetain<T> {
    fn retain_with_index<F>(&mut self, f: F)
    where
        F: FnMut(usize, &T) -> bool;
}

impl<T> IndexedRetain<T> for Vec<T> {
    fn retain_with_index<F>(&mut self, mut f: F)
    where
        F: FnMut(usize, &T) -> bool, // the signature of the callback changes
    {
        let len = self.len();
        let mut del = 0;
        {
            let v = &mut **self;

            for i in 0..len {
                // only implementation change here
                if !f(i, &v[i]) {
                    del += 1;
                } else if del > 0 {
                    v.swap(i - del, i);
                }
            }
        }
        if del > 0 {
            self.truncate(len - del);
        }
    }
}

such that the example would look like this:

v.retain_with_index(|index, item| (index % 2 == 0) || item == 4);

Or... better yet, you could use a higher-order function:

fn with_index<T, F>(mut f: F) -> impl FnMut(&T) -> bool
where
    F: FnMut(usize, &T) -> bool,
{
    let mut i = 0;
    move |item| (f(i, item), i += 1).0
}

such that the example would now look like this:

v.retain(with_index(|index, item| (index % 2 == 0) || item == 4));

(my preference is the latter)

Upvotes: 6

Related Questions