How can I retain vector elements with their original index?

Question

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]);

I don't see why. I have an algorithm that needs to know the index of items, and also wants to remove some of them. It can do it all in one pass. — Timmmm, Jan 05 '20 at 18:27
How can you possibly say that without knowing anything about it? — Timmmm, Jan 06 '20 at 09:48

score 6 · Accepted Answer · answered Jan 06 '20 at 18:49

@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)

You could also implement the trait function using the higher-order function if you really wanted to... — richardpringle, Jan 06 '20 at 18:54
I really how the solution at the bottom makes for a nice functional API. Could you expand on the implementation, in particular this line: `move |item| (f(i, item), i += 1).0` I'm new to Rust struggling to understand what's happening there. Thanks! — doplumi, Feb 01 '22 at 16:26
You can see that the index works as expected in the playground link I posted as part of the solution. What's happening in that closure is the `i` is moved by value to `f`. Since `i` is `Copy`, `f` is actually getting a copy of the value at the time it is passed into `f` as an argument. The result of that function call is then stored in the zeroth location of the tuple. After, `i` is incremented by 1, and the zeroth element is returned. A better way to understand how I wrote it is probably to check out line 40 in the playground. It's functionally the same. — richardpringle, Feb 02 '22 at 23:26

score 1 · Answer 2 · edited Jan 06 '20 at 18:52

1

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. :-)

edited Jan 06 '20 at 18:52

Shepmaster

388,571
95
1,107
1,366

answered Jan 05 '20 at 18:10

Timmmm

88,195
71
364
509

score 0 · Answer 3 · edited Jan 06 '20 at 18:52

0

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()

edited Jan 06 '20 at 18:52

Shepmaster

388,571
95
1,107
1,366

answered Jan 05 '20 at 17:19

Hauleth

22,873
4
61
112

Will that not allocate a new `Vec` rather than deleting from the existing one? – Timmmm Jan 05 '20 at 17:57
Yes, it will. However bench before discarding this solution. – Hauleth Jan 05 '20 at 19:32
2

I don't think benchmarking is necessary - this will definitely be slower. – Timmmm Jan 06 '20 at 09:42

How can I retain vector elements with their original index?

3 Answers3