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

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

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 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

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 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));
}