Skip to content

Instantly share code, notes, and snippets.

@sjkillen
Last active March 17, 2022 08:31
Show Gist options
  • Save sjkillen/3cfdd720cbc055a4e9692b127eb40a18 to your computer and use it in GitHub Desktop.
Save sjkillen/3cfdd720cbc055a4e9692b127eb40a18 to your computer and use it in GitHub Desktop.
Rust Cartesian Product Iterator
mod util;
use std::ops::Range;
use util::*;
fn main() {
let v: Vec<Range<i32>> = [(1..4), (11..14), (111..114)]
.into_iter()
.map(|r| r.clone())
.collect();
for items in v.into_iter().flat_cprod() {
println!("{:?}", items);
}
}
[1, 11, 111]
[1, 11, 112]
[1, 11, 113]
[1, 12, 111]
[1, 12, 112]
[1, 12, 113]
[1, 13, 111]
[1, 13, 112]
[1, 13, 113]
[2, 11, 111]
[2, 11, 112]
[2, 11, 113]
[2, 12, 111]
[2, 12, 112]
[2, 12, 113]
[2, 13, 111]
[2, 13, 112]
[2, 13, 113]
[3, 11, 111]
[3, 11, 112]
[3, 11, 113]
[3, 12, 111]
[3, 12, 112]
[3, 12, 113]
[3, 13, 111]
[3, 13, 112]
[3, 13, 113]
use std::iter::Peekable;
// Followed cartesian_product from Itertools
pub trait IteratorUtil: Iterator {
fn flat_cprod<J, I>(self) -> FlatCProd<J, I>
where
Self: Sized,
Self::Item: Clone + IntoIterator<Item = I, IntoIter = J>,
J: Clone + Iterator<Item = I>,
I: Clone,
{
let initial_columns: Vec<J> = self.map(|i| i.into_iter()).collect();
let columns = initial_columns
.clone()
.into_iter()
.map(|i| i.peekable())
.collect();
FlatCProd {
columns,
initial_columns,
}
}
}
#[derive(Debug, Clone)]
pub struct FlatCProd<J, I>
where
J: Clone + Iterator<Item = I>,
I: Clone,
{
initial_columns: Vec<J>,
columns: Vec<Peekable<J>>,
}
impl<J: Clone + Iterator<Item = I>, I: Clone> FlatCProd<J, I> {
fn peek_next(&mut self) -> Option<Vec<I>> {
return self
.columns
.iter_mut()
.map(|i| i.peek().map(|o| o.clone()))
.collect();
}
}
impl<J: Clone + Iterator<Item = I>, I: Clone> Iterator for FlatCProd<J, I> {
type Item = Vec<I>;
fn next(&mut self) -> Option<Self::Item> {
if self.columns.len() == 0 {
return None;
}
let next = self.peek_next();
match next {
None => return None,
Some(curr) => {
'l: for i in (0..self.columns.len()).rev() {
// This is always Some due to peek_next
self.columns[i].next();
match self.columns[i].peek() {
Some(_) => break 'l,
None => {
// Reset the column if it's not the last one
if i > 0 {
self.columns[i] = self.initial_columns[i].clone().peekable();
}
}
};
}
return Some(curr);
}
}
}
}
impl<I> IteratorUtil for I where I: Iterator {}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment