diff --git a/src/combinations.rs b/src/combinations.rs index e254fee7e..6bb2f3ec6 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -143,6 +143,27 @@ impl Combinations { // If we've made it this far, we haven't run out of combos false } + + /// Returns the n-th item or the number of successful steps. + pub(crate) fn try_nth(&mut self, n: usize) -> Result<::Item, usize> + where + I::Item: Clone, + { + let done = if self.first { + self.init() + } else { + self.increment_indices() + }; + if done { + return Err(0); + } + for i in 0..n { + if self.increment_indices() { + return Err(i + 1); + } + } + Ok(self.pool.get_at(&self.indices)) + } } impl Iterator for Combinations @@ -166,23 +187,7 @@ where } fn nth(&mut self, n: usize) -> Option { - let done = if self.first { - self.init() - } else { - self.increment_indices() - }; - - if done { - return None; - } - - for _ in 0..n { - if self.increment_indices() { - return None; - } - } - - Some(self.pool.get_at(&self.indices)) + self.try_nth(n).ok() } fn size_hint(&self) -> (usize, Option) { diff --git a/src/powerset.rs b/src/powerset.rs index 1b2b38540..734eaf614 100644 --- a/src/powerset.rs +++ b/src/powerset.rs @@ -42,6 +42,18 @@ where } } +impl Powerset { + /// Returns true if `k` has been incremented, false otherwise. + fn increment_k(&mut self) -> bool { + if self.combs.k() < self.combs.n() || self.combs.k() == 0 { + self.combs.reset(self.combs.k() + 1); + true + } else { + false + } + } +} + impl Iterator for Powerset where I: Iterator, @@ -52,14 +64,27 @@ where fn next(&mut self) -> Option { if let Some(elt) = self.combs.next() { Some(elt) - } else if self.combs.k() < self.combs.n() || self.combs.k() == 0 { - self.combs.reset(self.combs.k() + 1); + } else if self.increment_k() { self.combs.next() } else { None } } + fn nth(&mut self, mut n: usize) -> Option { + loop { + match self.combs.try_nth(n) { + Ok(item) => return Some(item), + Err(steps) => { + if !self.increment_k() { + return None; + } + n -= steps; + } + } + } + } + fn size_hint(&self) -> SizeHint { let k = self.combs.k(); // Total bounds for source iterator.