How to zip two iterators of unequal length with a default? - rust

Question

I'm trying to zip two iterators of unequal length, it only returns when when there is value in both and ignores the rest in the longest iterator.
fn main() {
let num1 = vec![1, 2];
let num2 = vec![3];
for i in num1.iter().rev().zip(num2.iter().rev()) {
println!("{:?}", i);
}
}
This returns (2, 3). How do i make it return:
(2, 3)
(1, 0) // default is the 0 here.
Is there any other way to do it?

Answer 1

You could use the zip_longest provided by the itertools crate.
use itertools::{
Itertools,
EitherOrBoth::*,
};
fn main() {
let num1 = vec![1, 2];
let num2 = vec![3];
for pair in num1.iter().rev().zip_longest(num2.iter().rev()) {
match pair {
Both(l, r) => println!("({:?}, {:?})", l, r),
Left(l) => println!("({:?}, 0)", l),
Right(r) => println!("(0, {:?})", r),
}
}
}
Which would produce the following output:
(2, 3)
(1, 0)

Answer 2

Zip will stop as soon as one of iterators stops producing values. If you know which is the longest, you can pad the shorter one with your default value:
use std::iter;
fn main() {
let longer = vec![1, 2];
let shorter = vec![3];
for i in longer
.iter()
.rev()
.zip(shorter.iter().rev().chain(iter::repeat(&0)))
{
println!("{:?}", i);
}
}
If you don't know which is longest, you should use itertools, as Peter Varo suggests.

Answer 3

The key is to detect that one iterator is shorter then the other, you could do it before before in your case vector implement ExactSizeIterator but a general solution would be to have a custom .zip().
itertools already offer a general solution, .zip_longest():
use itertools::EitherOrBoth::{Both, Left, Right};
use itertools::Itertools;
fn main() {
let num1 = vec![1, 2];
let num2 = vec![3];
for i in num1
.iter()
.rev()
.zip_longest(num2.iter().rev())
.map(|x| match x {
Both(a, b) => (a, b),
Left(a) => (a, &0),
Right(b) => (&0, b),
})
{
println!("{:?}", i);
}
}
This require you write the closure everytime, if you need this feature a lot maybe implement a custom trait on iterator with .zip_default() where A and B implement Default:
use std::default::Default;
use std::iter::Fuse;
pub trait MyIterTools: Iterator {
fn zip_default<J>(self, other: J) -> ZipDefault<Self, J::IntoIter>
where
J: IntoIterator,
Self: Sized,
{
ZipDefault::new(self, other.into_iter())
}
}
#[derive(Clone, Debug)]
pub struct ZipDefault<I, J> {
i: Fuse<I>,
j: Fuse<J>,
}
impl<I, J> ZipDefault<I, J>
where
I: Iterator,
J: Iterator,
{
fn new(i: I, j: J) -> Self {
Self {
i: i.fuse(),
j: j.fuse(),
}
}
}
impl<T, U, A, B> Iterator for ZipDefault<T, U>
where
T: Iterator<Item = A>,
U: Iterator<Item = B>,
A: Default,
B: Default,
{
type Item = (A, B);
fn next(&mut self) -> Option<Self::Item> {
match (self.i.next(), self.j.next()) {
(Some(a), Some(b)) => Some((a, b)),
(Some(a), None) => Some((a, B::default())),
(None, Some(b)) => Some((A::default(), b)),
(None, None) => None,
}
}
}
impl<T: ?Sized> MyIterTools for T where T: Iterator {}
fn main() {
let num1 = vec![1, 2];
let num2 = vec![3];
for i in num1
.iter()
.copied()
.rev()
.zip_default(num2.iter().copied().rev())
{
println!("{:?}", i);
}
}
Using itertools we can delegate some logic:
use std::default::Default;
use itertools::Itertools;
use itertools::ZipLongest;
use itertools::EitherOrBoth::{Both, Left, Right};
pub trait MyIterTools: Iterator {
fn zip_default<J>(self, j: J) -> ZipDefault<Self, J::IntoIter>
where
Self: Sized,
J: IntoIterator,
{
ZipDefault::new(self, j.into_iter())
}
}
#[derive(Clone, Debug)]
pub struct ZipDefault<I, J> {
inner: ZipLongest<I, J>,
}
impl<I, J> ZipDefault<I, J>
where
I: Iterator,
J: Iterator,
{
fn new(i: I, j: J) -> Self {
Self {
inner: i.zip_longest(j),
}
}
}
impl<T, U, A, B> Iterator for ZipDefault<T, U>
where
T: Iterator<Item = A>,
U: Iterator<Item = B>,
A: Default,
B: Default,
{
type Item = (A, B);
fn next(&mut self) -> Option<Self::Item> {
match self.inner.next()? {
Both(a, b) => Some((a, b)),
Left(a) => Some((a, B::default())),
Right(b) => Some((A::default(), b)),
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<T: ?Sized> MyIterTools for T where T: Iterator {}
fn main() {
let num1 = vec![1, 2];
let num2 = vec![3];
for i in num1
.iter()
.copied()
.rev()
.zip_default(num2.iter().copied().rev())
{
println!("{:?}", i);
}
}

Answer 4

I saw this neat trick in other guy code in leetcode solution. If you have access to length, you can swap iterators, making iter1 the longest.
fn main() {
let num1 = vec![1, 2];
let num2 = vec![3];
let mut iter1 = num1.iter();
let mut iter2 = num2.iter();
if iter1.len() < iter2.len(){
std::mem::swap(&mut iter1, &mut iter2);
} // now iter1 is the largest
for i in iter1.rev().zip(iter2.rev().chain(std::iter::repeat(&0))) {
println!("{:?}", i);
}
}

Answer 5

If you can get the length of the iterators, as is in this case, a quick and dirty way could be:
use std::iter::repeat;
fn main() {
let a = vec![1, 2, 3];
let b = vec![4, 5, 6, 7];
for i in a
.iter()
.rev()
.chain(repeat(&0).take(b.len().saturating_sub(a.len())))
.zip(
b.iter()
.rev()
.chain(repeat(&0).take(a.len().saturating_sub(b.len()))),
)
{
println!("{:?}", i);
}
}
You can also implement a trait containing a zip_default() using this approach:
pub trait MyIterTools<X: Default + Clone>: ExactSizeIterator<Item = X> {
fn zip_default<J, Y>(self, j: J) -> ZipDefault<Self, J::IntoIter, X, Y>
where
Self: Sized,
J: IntoIterator<Item = Y>,
J::IntoIter: ExactSizeIterator,
Y: Default + Clone,
{
ZipDefault::new(self, j.into_iter())
}
}
#[derive(Clone, Debug)]
pub struct ZipDefault<
I: ExactSizeIterator<Item = X>,
J: ExactSizeIterator<Item = Y>,
X: Default + Clone,
Y: Default + Clone,
> {
inner: Zip<Chain<I, Take<Repeat<X>>>, Chain<J, Take<Repeat<Y>>>>,
}
impl<
I: ExactSizeIterator<Item = X>,
J: ExactSizeIterator<Item = Y>,
X: Default + Clone,
Y: Default + Clone,
> ZipDefault<I, J, X, Y>
{
fn new(a: I, b: J) -> Self {
let a_len = a.len();
let b_len = b.len();
Self {
inner: a
.chain(repeat(X::default()).take(b_len.saturating_sub(a_len)))
.zip(b.chain(repeat(Y::default()).take(a_len.saturating_sub(b_len)))),
}
}
}
impl<
I: ExactSizeIterator<Item = X>,
J: ExactSizeIterator<Item = Y>,
X: Default + Clone,
Y: Default + Clone,
> Iterator for ZipDefault<I, J, X, Y>
{
type Item = (X, Y);
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<T: ExactSizeIterator<Item = X>, X: Default + Clone> MyIterTools<X> for T {}
fn main() {
let a = vec![1, 2, 3];
let b = vec![4, 5, 6, 7];
a.into_iter()
.zip_default(b.into_iter())
.for_each(|i| println!("{:?}", i));
}