@@ -1135,24 +1135,6 @@ The following recipes have a more mathematical flavor:
11351135 n = n // p * (p - 1)
11361136 return n
11371137
1138- def nth_combination(iterable, r, index):
1139- "Equivalent to list(combinations(iterable, r))[index]"
1140- pool = tuple(iterable)
1141- n = len(pool)
1142- c = math.comb(n, r)
1143- if index < 0:
1144- index += c
1145- if index < 0 or index >= c:
1146- raise IndexError
1147- result = []
1148- while r:
1149- c, n, r = c*r//n, n-1, r-1
1150- while index >= c:
1151- index -= c
1152- c, n = c*(n-r)//n, n-1
1153- result.append(pool[-1-n])
1154- return tuple(result)
1155-
11561138
11571139.. doctest ::
11581140 :hide:
@@ -1576,20 +1558,6 @@ The following recipes have a more mathematical flavor:
15761558 >>> first_true(' ABC0DEF1' , ' 9' , str .isdigit)
15771559 '0'
15781560
1579- >>> population = ' ABCDEFGH'
1580- >>> for r in range (len (population) + 1 ):
1581- ... seq = list (combinations(population, r))
1582- ... for i in range (len (seq)):
1583- ... assert nth_combination(population, r, i) == seq[i]
1584- ... for i in range (- len (seq), 0 ):
1585- ... assert nth_combination(population, r, i) == seq[i]
1586-
1587- >>> iterable = ' abcde'
1588- >>> r = 3
1589- >>> combos = list (combinations(iterable, r))
1590- >>> all (nth_combination(iterable, r, i) == comb for i, comb in enumerate (combos))
1591- True
1592-
15931561
15941562.. testcode ::
15951563 :hide:
@@ -1616,6 +1584,24 @@ The following recipes have a more mathematical flavor:
16161584 for (a, _), (b, c) in pairwise(pairwise(iterable)):
16171585 yield a, b, c
16181586
1587+ def nth_combination(iterable, r, index):
1588+ "Equivalent to list(combinations(iterable, r))[index]"
1589+ pool = tuple(iterable)
1590+ n = len(pool)
1591+ c = math.comb(n, r)
1592+ if index < 0:
1593+ index += c
1594+ if index < 0 or index >= c:
1595+ raise IndexError
1596+ result = []
1597+ while r:
1598+ c, n, r = c*r//n, n-1, r-1
1599+ while index >= c:
1600+ index -= c
1601+ c, n = c*(n-r)//n, n-1
1602+ result.append(pool[-1-n])
1603+ return tuple(result)
1604+
16191605
16201606.. doctest ::
16211607 :hide:
@@ -1631,3 +1617,17 @@ The following recipes have a more mathematical flavor:
16311617
16321618 >>> list (triplewise(' ABCDEFG' ))
16331619 [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E'), ('D', 'E', 'F'), ('E', 'F', 'G')]
1620+
1621+ >>> population = ' ABCDEFGH'
1622+ >>> for r in range (len (population) + 1 ):
1623+ ... seq = list (combinations(population, r))
1624+ ... for i in range (len (seq)):
1625+ ... assert nth_combination(population, r, i) == seq[i]
1626+ ... for i in range (- len (seq), 0 ):
1627+ ... assert nth_combination(population, r, i) == seq[i]
1628+
1629+ >>> iterable = ' abcde'
1630+ >>> r = 3
1631+ >>> combos = list (combinations(iterable, r))
1632+ >>> all (nth_combination(iterable, r, i) == comb for i, comb in enumerate (combos))
1633+ True
0 commit comments