Add the nth_combination itertools recipe as a test case (#11031)

This commit is contained in:
Alex Waygood
2023-11-15 19:14:21 +00:00
committed by GitHub
parent ac88fdfdbf
commit 4aae235cef

View File

@@ -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