Implement Combinations::nth

src/combinations.rs

@@ -94,6 +94,55 @@ impl<I: Iterator> Combinations<I> {
         let n = pool.count();
         (n, remaining_for(n, first, &indices).unwrap())
     }
+
+    /// Initialises the iterator by filling a buffer with elements from the
+    /// iterator. Returns true if there are no combinations, false otherwise.
+    fn init(&mut self) -> bool {
+        self.pool.prefill(self.k());
+        let done = self.k() > self.n();
+        if !done {
+            self.first = false;
+        }
+
+        done
+    }
+
+    /// Increments indices representing the combination to advance to the next
+    /// (in lexicographic order by increasing sequence) combination. For example
+    /// if we have n=4 & k=2 then `[0, 1] -> [0, 2] -> [0, 3] -> [1, 2] -> ...`
+    ///
+    /// Returns true if we've run out of combinations, false otherwise.
+    fn increment_indices(&mut self) -> bool {
+        if self.indices.is_empty() {
+            return true; // Done
+        }
+
+        // Scan from the end, looking for an index to increment
+        let mut i: usize = self.indices.len() - 1;
+
+        // Check if we need to consume more from the iterator
+        if self.indices[i] == self.pool.len() - 1 {
+            self.pool.get_next(); // may change pool size
+        }
+
+        while self.indices[i] == i + self.pool.len() - self.indices.len() {
+            if i > 0 {
+                i -= 1;
+            } else {
+                // Reached the last combination
+                return true;
+            }
+        }
+
+        // Increment index, and reset the ones to its right
+        self.indices[i] += 1;
+        for j in i + 1..self.indices.len() {
+            self.indices[j] = self.indices[j - 1] + 1;
+        }
+
+        // If we've made it this far, we haven't run out of combos
+        false
+    }
 }
 
 impl<I> Iterator for Combinations<I>
@@ -103,40 +152,40 @@ where
 {
     type Item = Vec<I::Item>;
     fn next(&mut self) -> Option<Self::Item> {
-        if self.first {
-            self.pool.prefill(self.k());
-            if self.k() > self.n() {
-                return None;
-            }
-            self.first = false;
-        } else if self.indices.is_empty() {
-            return None;
+        let done = if self.first {
+            self.init()
         } else {
-            // Scan from the end, looking for an index to increment
-            let mut i: usize = self.indices.len() - 1;
+            self.increment_indices()
+        };
 
-            // Check if we need to consume more from the iterator
-            if self.indices[i] == self.pool.len() - 1 {
-                self.pool.get_next(); // may change pool size
-            }
+        if done {
+            return None;
+        }
 
-            while self.indices[i] == i + self.pool.len() - self.indices.len() {
-                if i > 0 {
-                    i -= 1;
-                } else {
-                    // Reached the last combination
-                    return None;
-                }
-            }
+        Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect())
+    }
 
-            // Increment index, and reset the ones to its right
-            self.indices[i] += 1;
-            for j in i + 1..self.indices.len() {
-                self.indices[j] = self.indices[j - 1] + 1;
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        if n == 0 {
+            return self.next();
+        }
+
+        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;
             }
         }
 
-        // Create result vector based on the indices
         Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect())
     }