aoc/src/util/iter.rs
2023-11-26 23:15:08 +01:00

102 lines
2.0 KiB
Rust

pub struct Combinations<I, T>
where
I: Iterator<Item = T>,
{
indices: Vec<usize>,
k: usize,
size: usize,
pool: Vec<I::Item>,
first: bool,
}
impl<I, T> Combinations<I, T>
where
I: Iterator<Item = T>,
{
pub fn new(iter: I, k: usize) -> Self {
let pool: Vec<I::Item> = iter.collect();
if k > pool.len() {
panic!(
"No combinations possible for len {} and k {}",
pool.len(),
k
);
}
Self {
indices: (0..k).collect(),
k,
size: pool.len(),
pool,
first: true,
}
}
}
impl<I, T> Iterator for Combinations<I, T>
where
I: Iterator<Item = T>,
T: Copy,
{
type Item = Vec<T>;
fn next(&mut self) -> Option<Self::Item> {
if self.first {
self.first = false;
} else {
let mut i = self.k - 1;
while self.indices[i] == i + self.size - self.k {
if i > 0 {
i -= 1;
} else {
return None;
}
}
self.indices[i] += 1;
for j in i + 1..self.k {
self.indices[j] = self.indices[j - 1] + 1;
}
}
Some(self.indices.iter().map(|i| self.pool[*i]).collect())
}
}
pub trait CombinationsIterator<T>: Iterator<Item = T> + Sized {
fn combinations(self, k: usize) -> Combinations<Self, T> {
Combinations::new(self, k)
}
}
impl<I, T> CombinationsIterator<T> for I where I: Iterator<Item = T> {}
pub struct Window<const SIZE: usize, I, T>
where
I: Iterator<Item = T>,
{
iter: I,
}
impl<const SIZE: usize, I, T> Window<SIZE, I, T>
where
I: Iterator<Item = T>,
{
pub fn new(iter: I) -> Self {
Self { iter }
}
}
impl<const SIZE: usize, I, T> Iterator for Window<SIZE, I, T>
where
I: Iterator<Item = T>,
{
type Item = [I::Item; SIZE];
fn next(&mut self) -> Option<Self::Item> {
todo!()
}
}