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.