[Rust] イテレータで直積集合を生成

適当に作ったので間違ってると思います。適宜修正して使ってください。

注意 一番簡単かつ普通なのはこちらの方法です。本文中のコードはかなり変態なので、これを見てRustのことを嫌いにならないでください。

モチベーション

let v = vec![1, 2];
let w = vec![3, 4];

という2つの集合が与えられた時、

(1, 3)
(1, 4)
(2, 3)
(2, 4)

というペアのイテレータを取得したい。

Code

struct Cross<T: Iterator, U: Iterator + Clone>(T, U, Option<T::Item>, Option<U>)
where
    T::Item: Clone;

fn cross<T, U>(this: T, other: U) -> Cross<T::IntoIter, U::IntoIter>
where
    T: IntoIterator,
    T::Item: Clone,
    U: IntoIterator,
    U::IntoIter: Clone,
{
    Cross(this.into_iter(), other.into_iter(), None, None)
}
impl<T, U> Iterator for Cross<T, U>
where
    T: Iterator,
    T::Item: Clone,
    U: Iterator + Clone,
{
    type Item = (T::Item, U::Item);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if let None = self.2 {
            if let Some(v1) = self.0.next() {
                self.2 = Some(v1);
            } else {
                return None;
            }
        }
        if let None = self.3 {
            self.3 = Some(self.1.clone())
        }
        if let Some(v2) = self.3.as_mut().unwrap().next() {
            Some((self.2.as_ref().unwrap().clone(), v2))
        } else {
            self.2 = None;
            self.3 = None;
            self.next()
        }
    }
}
fn main() {
    let v = vec![1, 2, 3];
    let w = vec![4, 5, 6];
    for (i, j) in cross(v, w) {
        println!("({}, {})", i, j);
    }
}

(追加) 別のバージョン

標準ライブラリを用いて手を抜いた...つもりが前のバージョンより時間がかかりました

use std::iter::{Cycle, FlatMap, Map, Zip};
use std::vec::IntoIter;
fn product<S, T>(
    ls: S,
    rs: T,
) -> FlatMap<
    Zip<S::IntoIter, Cycle<IntoIter<T>>>,
    Map<Zip<T::IntoIter, Cycle<IntoIter<S::Item>>>, fn((T::Item, S::Item)) -> (S::Item, T::Item)>,
    fn(
        (S::Item, T),
    ) -> Map<
        Zip<T::IntoIter, Cycle<IntoIter<S::Item>>>,
        fn((T::Item, S::Item)) -> (S::Item, T::Item),
    >,
>
where
    S: IntoIterator,
    T: IntoIterator + Clone,
    S::Item: Clone,
{
    ls.into_iter()
        .zip(vec![rs].into_iter().cycle())
        .flat_map(|(l, rs)| {
            rs.into_iter()
                .zip(vec![l].into_iter().cycle())
                .map(|(r, l)| (l, r))
        })
}

fn main() {
    let xs = vec![1, 2, 3];
    let ys = vec![4, 5, 6];

    println!("{:?}", product(xs, ys).collect::<Vec<_>>());
}

外部の変数をキャプチャするクロージャの型修飾をうまく書けないので、無理やりzipとcycleを用いてキャプチャする変数を引数に持っていった。

タイプエイリアスを使ったら簡単になるかなと思ったけど大して変わらなかった

use std::iter::{Cycle, FlatMap, Map, Zip};
use std::vec::IntoIter;
type M<S, T> = Map<
    Zip<<T as IntoIterator>::IntoIter, Cycle<IntoIter<<S as IntoIterator>::Item>>>,
    fn(
        (<T as IntoIterator>::Item, <S as IntoIterator>::Item),
    ) -> (<S as IntoIterator>::Item, <T as IntoIterator>::Item),
>;
fn product<S, T>(
    ls: S,
    rs: T,
) -> FlatMap<Zip<S::IntoIter, Cycle<IntoIter<T>>>, M<S, T>, fn((S::Item, T)) -> M<S, T>>
where
    S: IntoIterator,
    T: IntoIterator + Clone,
    S::Item: Clone,
{
    ls.into_iter()
        .zip(vec![rs].into_iter().cycle())
        .flat_map(|(l, rs)| {
            rs.into_iter()
                .zip(vec![l].into_iter().cycle())
                .map(|(r, l)| (l, r))
        })
}

fn main() {
    let xs = vec![1, 2, 3];
    let ys = vec![4, 5, 6];

    println!("{:?}", product(xs, ys).collect::<Vec<_>>());
}

これはImpl Traitを用いるとさらに簡単になります。

(追加) もう少し理解しやすいバージョン

内側のループでイテレータではなくVecを生成するようにした。

use std::iter::{Cycle, FlatMap, Zip};
use std::vec::IntoIter;
fn product<S, T>(
    ls: S,
    rs: T,
) -> FlatMap<
    Zip<S::IntoIter, Cycle<IntoIter<T>>>,
    Vec<(S::Item, T::Item)>,
    fn((S::Item, T)) -> Vec<(S::Item, T::Item)>,
>
where
    S: IntoIterator,
    T: IntoIterator + Clone,
    S::Item: Clone,
{
    ls.into_iter()
        .zip(vec![rs].into_iter().cycle())
        .flat_map(|(l, rs)| rs.into_iter().map(|r| (l.clone(), r)).collect::<Vec<_>>())
}

fn main() {
    let xs = vec![1, 2, 3];
    let ys = vec![4, 5, 6];

    println!("{:?}", product(xs, ys).collect::<Vec<_>>());
}

おわりに

標準ライブラリでやる方法があったら教えてください

#rustlang の #iterator で直積集合をうまく作る方法ないかな?
今の所、fold()で適当にやってるけど

— yasuo_ozu (@yasuo_ozu) September 12, 2020