CODEWITHSUNDEEP
rustiterator
picked name
picked name

Reputation: 399

Is there a way to merge two iterators into the same Vec?

I have two iterators that will end up merging into the same Vec, but I need to perform a filter on the end result before merging. For example:

let a = vec![1, 2, 3, 4].into_iter().map(|x| x * 2);
let b = vec![0, 3, 5, 6, 7].into_iter().map(|x| x * 3);
let c = a + b;
assert_eq!(
    c.filter(|&x| x > 5).collect::<Vec<u8>>(),
    vec![6, 8, 9, 15, 18, 21]
);

I could do something like this:

let mut a = vec![1, 2, 3, 4]
    .into_iter()
    .map(|x| x * 2)
    .collect::<Vec<u8>>();
let b = vec![0, 3, 5, 6, 7]
    .into_iter()
    .map(|x| x * 3)
    .collect::<Vec<u8>>();
a.extend(b);
assert_eq!(
    a.into_iter().filter(|&x| x > 5).collect::<Vec<u8>>(),
    vec![6, 8, 9, 15, 18, 21]
);

But the extra allocation kills performance in my case (yes, I checked!)

Upvotes: 12

Views: 16250

Answers (2)

Stargateur
Stargateur

Reputation: 26767

Alternatively, you could use extend():

let a = (0..500000).map(|x| x * 2);
let b = (0..500000).map(|x| x * 3);
let mut c = Vec::with_capacity(a.size_hint().1.unwrap() + b.size_hint().1.unwrap());
c.extend(a);
c.extend(b);

This requires you to explicitly use with_capacity() which chain() and collect() would have done for you. In the situation where a is already constructed, extend() will be suitable instead of building a temporary vector.

I did not find any benchmark difference between chain() and extend in this case (benchmark)

This doesn't do the filtering. Sebastian Redl

Correct! Fixing this mistake shows that for some reason, LLVM no longer optimizes chain() now. The version with extend:

let a = (0..500000).map(|x| x * 2);
let b = (0..500000).map(|x| x * 3);
let mut c = Vec::with_capacity(a.size_hint().1.unwrap() + b.size_hint().1.unwrap());
c.extend(a.filter(|&x| x > 5));
c.extend(b.filter(|&x| x > 5));

Is twice as fast than the version with chain (this is what I expected in the first place):

let a = (0..500000).map(|x| x * 2);
let b = (0..500000).map(|x| x * 3);
let _c: Vec<_> = a.chain(b).filter(|&x| x > 5).collect();

(benchmark)

Upvotes: 1

Sebastian Redl
Sebastian Redl

Reputation: 72063

You're looking for Iterator::chain

let a = vec![1, 2, 3, 4].into_iter().map(|x| x * 2);
let b = vec![0, 3, 5, 6, 7].into_iter().map(|x| x * 3);
let c = a.chain(b);
assert_eq!(
    c.filter(|&x| x > 5).collect::<Vec<u8>>(),
    vec![6, 8, 9, 15, 18, 21]
);

Upvotes: 24

Related Questions