From dd6a56932141c8b18b9d479eba30134c5afbfb88 Mon Sep 17 00:00:00 2001
From: Philippe-Cholet <44676486+Philippe-Cholet@users.noreply.github.com>
Date: Thu, 25 Apr 2024 11:15:56 +0200
Subject: [PATCH] Specialize `CombinationsWithReplacement::nth`

Similar to `next` except we increment indices n times before generating the vector item.
---
 src/combinations_with_replacement.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs
index 7acfee89..f363f9ba 100644
--- a/src/combinations_with_replacement.rs
+++ b/src/combinations_with_replacement.rs
@@ -106,6 +106,24 @@ where
         Some(self.pool.get_at(&self.indices))
     }
 
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        if self.first {
+            // In empty edge cases, stop iterating immediately
+            if !(self.indices.is_empty() || self.pool.get_next()) {
+                return None;
+            }
+            self.first = false;
+        } else if self.increment_indices() {
+            return None;
+        }
+        for _ in 0..n {
+            if self.increment_indices() {
+                return None;
+            }
+        }
+        Some(self.pool.get_at(&self.indices))
+    }
+
     fn size_hint(&self) -> (usize, Option<usize>) {
         let (mut low, mut upp) = self.pool.size_hint();
         low = remaining_for(low, self.first, &self.indices).unwrap_or(usize::MAX);