Specialize Powerset::nth

rust-itertools · May 7, 2024 · 0f3cbcc · 0f3cbcc
1 parent 54429a0
commit 0f3cbcc
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/powerset.rs b/src/powerset.rs
@@ -71,6 +71,20 @@ where
         }
     }
 
+    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.