Remove a sequence of values from a Vec in Rust?

Question

remove on a Vec removes a single value, given by an index, and returns that value.

I want to remove a list of indices, for example remove indices 1, 2 and 5 from a Vec of length 8, and get the values at those indices as another Vec. Calling remove repeatedly is (a) expensive, and (b) error-prone, as after each remove other indices are moved.

So, if I started with let mut v = vec![2,3,4,5,6,7], and removed indices [1,2,5], I would end up with a new vector containing vec![3,4,7], while v would be vec![2,5,6].

Answer 1

I was just solving this for my own application, so I'm gonna share the module. It's not perfect and I'm sure it can be optimized, but it works well enough for me.

take_multiple() or take_multiple_in_order() will do exactly what you wanted in the original question.

pub trait RemoveMultiple<T> {
    /// Remove multiple indices
    fn remove_multiple(&mut self, to_remove: Vec<usize>);
    
    /// Remove multiple indices with swap_remove, this is faster but reorders elements
    fn swap_remove_multiple(&mut self, to_remove: Vec<usize>);
    
    /// Remove and return multiple indices
    fn take_multiple(&mut self, to_remove: Vec<usize>) -> Vec<T>;
        
    /// Remove and return multiple indices, preserving the order specified in the index list
    fn take_multiple_in_order(&mut self, to_remove: &[usize]) -> Vec<T>;
    
    /// Remove and return multiple indices with swap_remove, this is faster but reorders elements and the results are in reverse order
    fn swap_take_multiple(&mut self, to_remove: Vec<usize>) -> Vec<T>;
}

impl<T> RemoveMultiple<T> for Vec<T> {
    fn remove_multiple(&mut self, mut to_remove: Vec<usize>) {
        to_remove.sort();
        to_remove.reverse();
        for r in to_remove {
            self.remove(r);
        }
    }
    
    fn swap_remove_multiple(&mut self, mut to_remove: Vec<usize>) {
        to_remove.sort();
        to_remove.reverse();
        for r in to_remove {
            self.swap_remove(r);
        }
    }
    
    fn take_multiple(&mut self, mut to_remove: Vec<usize>) -> Vec<T> {
        to_remove.sort();
        to_remove.reverse();
        let mut collected = vec![];
        for r in to_remove {
            collected.push(self.remove(r));
        }
        collected.reverse();
        collected
    }
    
    fn take_multiple_in_order(&mut self, to_remove: &[usize]) -> Vec<T> {
        let mut to_remove = to_remove.iter().copied().enumerate().collect::<Vec<_>>();
        to_remove.sort_by_key(|(_, r)| *r);
        to_remove.reverse();
        let mut collected : Vec<Option<T>> = std::iter::repeat_with(|| None).take(to_remove.len()).collect();
        for (i, r) in to_remove {
            collected[i] = Some(self.remove(r));
        }
        collected.into_iter().filter_map(|x| x).collect()
    }
    
    fn swap_take_multiple(&mut self, mut to_remove: Vec<usize>) -> Vec<T> {
        to_remove.sort();
        to_remove.reverse();
        let mut collected = vec![];
        for r in to_remove {
            collected.push(self.swap_remove(r));
        }
        collected
    }
}


#[cfg(test)]
mod test {
    use super::*;
    
    #[test]
    fn remove_multiple() {
        let mut list = vec!['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];        
        list.remove_multiple(vec![8, 0, 5, 6]);
        assert_eq!(vec!['1', '2', '3', '4','7', '9'], list);             
    }
    
    #[test]
    fn swap_remove_multiple() {
        let mut list = vec!['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];        
        list.swap_remove_multiple(vec![8, 0, 5, 6]);
        assert_eq!(vec!['9', '1', '2', '3', '4', '7'], list);      
    }
    
    #[test]
    fn take_multiple() {
        let mut list = vec!['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];        
        let taken = list.take_multiple(vec![8, 0, 5, 6]);
        assert_eq!(vec!['1', '2', '3', '4','7', '9'], list); 
        assert_eq!(vec!['0', '5', '6', '8'], taken);             
    }
    
    #[test]
    fn swap_take_multiple() {
        let mut list = vec!['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];        
        let taken = list.swap_take_multiple(vec![8, 0, 5, 6]);
        assert_eq!(vec!['9', '1', '2', '3', '4', '7'], list); 
        assert_eq!(vec!['8', '6', '5', '0'], taken);                
    }
    
    #[test]
    fn take_multiple_in_order() {
        let mut list = vec!['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];        
        let taken = list.take_multiple_in_order(&vec![8, 0, 5, 6]);
        assert_eq!(vec!['1', '2', '3', '4','7', '9'], list); 
        assert_eq!(vec!['8', '0', '5', '6'], taken);             
    }
}

Answer 2

If the indices you want to remove are contiguous, you can use Vec::drain, as shown in this answer. If not (i.e. you have arbitrary indices you want to remove), things get a lot more complicated.

One can solve the "expensive" problem of remove by using swap_remove. Instead of shifting all elements to the left, it will swap the removed element with the last one, thus being a O(1) operation. However, with this it's still very error-prone as indices of elements change after each remove-operation. Additionally, the order of elements is not the same as before which might not work for you.

The only way (I can think of) to efficiently remove multiple arbitrary indices is to sort these indices in decreasing order.

/// `indices_to_remove` have to be sorted in decreased order!
fn remove_multiple<T>(source: &mut Vec<T>, indices_to_remove: &[usize]) -> Vec<T> {
    indices_to_remove.iter()
        .copied()
        .map(|i| source.swap_remove(i))
        .collect()
}

Example (Playground):

let mut vec = vec!['a', 'b', 'c', 'd', 'e', 'f'];
let removed = remove_multiple(&mut vec, &[5, 3, 2, 0]);
println!("source: {:?}", vec);        // ['e', 'b']
println!("result: {:?}", removed);    // ['f', 'd', 'c', 'a']

If your list of indices is not sorted, I actually think just sorting it is the most efficient way to achieve your goal. At least I can't think of an algorithm that keeps track of all the indices and that is faster than O(n * log n), which is the runtime of sorting first.