From 4aae235cefd62ec417b92166d2bcf038cb8801e8 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 15 Nov 2023 19:14:21 +0000
Subject: [PATCH] Add the `nth_combination` itertools recipe as a test case
 (#11031)

---
 .../itertools/check_itertools_recipes.py      | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/test_cases/stdlib/itertools/check_itertools_recipes.py b/test_cases/stdlib/itertools/check_itertools_recipes.py
index e90b76c34..b74d92ea5 100644
--- a/test_cases/stdlib/itertools/check_itertools_recipes.py
+++ b/test_cases/stdlib/itertools/check_itertools_recipes.py
@@ -289,6 +289,27 @@ def polynomial_derivative(coefficients: Sequence[float]) -> list[float]:
     return list(map(operator.mul, coefficients, powers))
 
 
+if sys.version_info >= (3, 8):
+
+    def nth_combination(iterable: Iterable[_T], r: int, index: int) -> tuple[_T, ...]:
+        "Equivalent to list(combinations(iterable, r))[index]"
+        pool = tuple(iterable)
+        n = len(pool)
+        c = math.comb(n, r)
+        if index < 0:
+            index += c
+        if index < 0 or index >= c:
+            raise IndexError
+        result: list[_T] = []
+        while r:
+            c, n, r = c * r // n, n - 1, r - 1
+            while index >= c:
+                index -= c
+                c, n = c * (n - r) // n, n - 1
+            result.append(pool[-1 - n])
+        return tuple(result)
+
+
 if sys.version_info >= (3, 10):
 
     @overload