Collect vec elements into multiple groups

#1

Starting from this code (Playground link)

enum Kind {
    Foo,
    Bar,
    Baz,
}

struct Element {
    kind: Kind,
    value: i32,
}

fn main() {
    let v = vec![
        Element { kind: Kind::Foo, value: 1 },
        Element { kind: Kind::Bar, value: 2 },
        Element { kind: Kind::Baz, value: 3 },
        Element { kind: Kind::Bar, value: 4 },
        Element { kind: Kind::Baz, value: 5 },
    ];
}

I'd like to filter and group the elements of v by their kind, discarding some types. So in the end I'd like to get two Vecs, looking like this:

assert!(foos, vec![
  Element { kind: Kind::Foo, value: 1 },
]);

assert!(bars, vec![
  Element { kind: Kind::Bar, value: 2 },
  Element { kind: Kind::Bar, value: 4 },
]);

I fiddled around with unzip():

let (foos, bars): (Vec<_>, Vec<_>) v.into_iter().map(|el| {
  match el.kind {
    Kind::Foo => (el, None),
    Kind::Foo => (None, el),
    _ => (None, None),
  }
}).unzip();

but I would still have to get rid of the Options.

#2

If you just need to split into two lists, then Iterator::partition() should do the trick:

    let (foos, bars): (Vec<_>, Vec<_>) = v.iter()
        .filter(|x| x.kind == Kind::Foo || x.kind == Kind::Bar)
        .partition(|x| x.kind == Kind::Foo);

(Playground link)

For brevity I'm assuming that #[derive(PartialEq, Eq]) has been added to the Kind enum. Otherwise, you can use match, which would be a bit longer.

1 Like
#3

You might be interested in Itertools::partition_map, perhaps with a flat_map first to filter things.

#4

For arbitrary grouping, you can also use a HashMap or a BTreeMap of Vecs.

#5

I would use a for loop: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=701b435ac5e1609be13bfcc2f7f69906

2 Likes