Specialize Powerset::nth

src/combinations.rs

@@ -143,6 +143,27 @@ impl<I: Iterator> Combinations<I> {
  // 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<<Self as Iterator>::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<I> Iterator for Combinations<I>
@@ -166,23 +187,7 @@ where
  }
 
  fn nth(&mut self, n: usize) -> Option<Self::Item> {
- 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<usize>) {

src/powerset.rs

@@ -42,6 +42,18 @@ where
  }
 }
 
+impl<I: Iterator> Powerset<I> {
+ /// 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<I> Iterator for Powerset<I>
 where
  I: Iterator,
@@ -52,14 +64,27 @@ where
  fn next(&mut self) -> Option<Self::Item> {
  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<Self::Item> {
+ 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.