Boolean masks for vectors (or iterators)

user_18811786 0
Viewed 131

What's the most efficient way to implement boolean masking with vectors in Rust?

There might be a crate that does this(?), but to help me learn, I'd also like to hear about:

  • how you'd do this without a crate (i.e. have I missed something in the standard library?), and;
  • whether bool masking can be extended to lists and iterators more generally.

My example below applies the bool vector [true, false, false, true] to mask out the true indices of a u32 vector [1, 2, 3, 4]. This returns [2, 3].

I'm using a clunky for loop; can we do it better using e.g. maps and iterators?

fn main() {
    let mut numbers = vec![1, 2, 3, 4];
    let mask = vec![true, false, false, true];

    remove_true(&mut numbers, &mask); // remove numbers where vector mask == true
    println!("{:?}", numbers)
}

// Using a for loop
fn remove_true(vec: &mut Vec<u32>, locs: &Vec<bool>) {
    for m in (0..locs.len()).rev() {
        // Loop in reverse to preserve indices when dropping values
        if locs[m] == true {
            vec.remove(m);
        }
    }
}
rust
2 Answers

Vec::retain can be used for this operation, in fact this sort of selection is one of the example snippets:

let mut vec = vec![1, 2, 3, 4, 5];
let keep = [false, true, true, false, true];
let mut iter = keep.iter();
vec.retain(|_| *iter.next().unwrap());
assert_eq!(vec, [2, 3, 5]);

Since you're performing a negative selection (excluding the selected items) you have to reverse the result, but that aside it's identical (in fact it's a tad simpler because ! doesn't require dereferencing):

fn remove_true(vec: &mut Vec<u32>, locs: &[bool]) {
    let mut drop = locs.into_iter();
    vec.retain(|_| !drop.next().unwrap())
}
user_81821180

A masking can be abstracted as a zip+filter over iterators too:

fn main() {
    let numbers = vec![1, 2, 3, 4];
    let mask = vec![true, false, false, true];

    let filtered: Vec<_> = numbers
        .iter()
        .zip(mask.iter())
        .filter_map(|(v, c)| if c { Some(v) } else { None })
        .collect();
    println!("{:?}", numbers)
}
user_16951720
Related