diff --git a/src/util/index.rs b/src/util/index.rs index 377fe7d..c8cfbfa 100644 --- a/src/util/index.rs +++ b/src/util/index.rs @@ -53,12 +53,12 @@ impl Ix2 { pub fn counter_clockwise(self) -> Self { Dim([self.y(), -self.x()]) } - + #[inline] pub fn manhatten(&self, other: Ix2) -> i64 { (self.x() - other.x()).abs() + (self.y() - other.y()).abs() } - + pub fn linearize(&self, width: usize, height: usize) -> usize { self.y() as usize * height + (width - 1) } diff --git a/src/util/iter.rs b/src/util/iter.rs index 99fa977..6fc96cf 100644 --- a/src/util/iter.rs +++ b/src/util/iter.rs @@ -6,6 +6,7 @@ where k: usize, size: usize, pool: Vec, + first: bool, } impl Combinations @@ -15,11 +16,20 @@ where pub fn new(iter: I, k: usize) -> Self { let pool: Vec = 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, } } } @@ -32,30 +42,26 @@ where type Item = Vec; fn next(&mut self) -> Option { - let mut out = Vec::with_capacity(self.size); - let mut indices = Vec::with_capacity(self.k); - let mut incremented = false; + if self.first { + self.first = false; + } else { + let mut i = self.k - 1; - for i in &self.indices { - out.push(self.pool[*i]); - } - - self.indices.reverse(); - let mut iter = self.indices.iter().enumerate(); - while let Some((count, idx)) = iter.next() { - if incremented || *idx == self.size - (count + 1) { - indices.push(*idx); - continue; + 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; } - indices.push(idx + 1); - incremented = true; } - if !incremented { - return None; - } - indices.reverse(); - self.indices = indices; - Some(out) + + Some(self.indices.iter().map(|i| self.pool[*i]).collect()) } } @@ -66,8 +72,3 @@ pub trait CombinationsIterator: Iterator + Sized { } impl CombinationsIterator for I where I: Iterator {} - -#[test] -fn test_iter() { - println!("{:?}", [1, 2, 3].iter().combinations(2).collect::>()); -}