Filter a vector using a boolean array

user_8053020 0
Viewed 237

How do I use the boolean_arr to filter ys based on the bool values?

The arrays all have the same length.

fn main() {
    let xs: Vec<f32> = vec![300., 7.5, 10., 250.];
    let boolean_arr: Vec<bool> = xs.into_iter().map(|x| x > 10.).collect();

    let ys: Vec<f32> = vec![110.5, 50., 25., 770.];

    assert_eq!(wanted_vec, vec![110.5, 770.]);
}
rust
2 Answers

We can use zip to combine the two and then filter the results. Here's a general-purpose function that works on arbitrary iterators.

fn filter_by<I1, I2>(bs: I1, ys: I2) -> impl Iterator<Item=<I2 as Iterator>::Item>
where I1: Iterator<Item=bool>,
      I2: Iterator {
  bs.zip(ys).filter(|x| x.0).map(|x| x.1)
}

And how to call it in your particular use case

let wanted_vec: Vec<_> = filter_by(boolean_arr.into_iter(), ys.into_iter()).collect();
user_22886590

A different tack from Silvio's answer: a more efficient (but less generic) method would be to use Vec::retain since it's guaranteed to traverse the subject in-order and you're guaranteeing that the subject and filter bitmap are the same size:

// don't even collect if not necessary, though that also applies to Silvio's suggestion
let mut bitmap = xs.into_iter().map(|x| x > 10.);

let mut ys: Vec<f32> = vec![110.5, 50., 25., 770.];
// could even use `unwrap_unchecked` if that's better codegen, but then that's UB if the bitmap length doesn't match the subject
ys.retain(|_| bitmap.next().unwrap());
let wanted_vec = ys;
user_81821180
Related