Skip to content

Commit d37258d

Browse files
authored
Add itertools recipe for directly finding the n-th combination (#5161)
1 parent 0f31c74 commit d37258d

File tree

2 files changed

+56
-0
lines changed

2 files changed

+56
-0
lines changed

Doc/library/itertools.rst

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -859,6 +859,29 @@ which incur interpreter overhead.
859859
indices = sorted(random.randrange(n) for i in range(r))
860860
return tuple(pool[i] for i in indices)
861861

862+
def nth_combination(iterable, r, index):
863+
'Equivalent to list(combinations(iterable, r))[index]'
864+
pool = tuple(iterable)
865+
n = len(pool)
866+
if r < 0 or r > n:
867+
raise ValueError
868+
c = 1
869+
k = min(r, n-r)
870+
for i in range(1, k+1):
871+
c = c * (n - k + i) // i
872+
if index < 0:
873+
index += c
874+
if index < 0 or index >= c:
875+
raise IndexError
876+
result = []
877+
while r:
878+
c, n, r = c*r//n, n-1, r-1
879+
while index >= c:
880+
index -= c
881+
c, n = c*(n-r)//n, n-1
882+
result.append(pool[-1-n])
883+
return tuple(result)
884+
862885
Note, many of the above recipes can be optimized by replacing global lookups
863886
with local variables defined as default values. For example, the
864887
*dotproduct* recipe can be written as::

Lib/test/test_itertools.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2262,6 +2262,30 @@ def test_permutations_sizeof(self):
22622262
... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
22632263
... return next(filter(pred, iterable), default)
22642264
2265+
>>> def nth_combination(iterable, r, index):
2266+
... 'Equivalent to list(combinations(iterable, r))[index]'
2267+
... pool = tuple(iterable)
2268+
... n = len(pool)
2269+
... if r < 0 or r > n:
2270+
... raise ValueError
2271+
... c = 1
2272+
... k = min(r, n-r)
2273+
... for i in range(1, k+1):
2274+
... c = c * (n - k + i) // i
2275+
... if index < 0:
2276+
... index += c
2277+
... if index < 0 or index >= c:
2278+
... raise IndexError
2279+
... result = []
2280+
... while r:
2281+
... c, n, r = c*r//n, n-1, r-1
2282+
... while index >= c:
2283+
... index -= c
2284+
... c, n = c*(n-r)//n, n-1
2285+
... result.append(pool[-1-n])
2286+
... return tuple(result)
2287+
2288+
22652289
This is not part of the examples but it tests to make sure the definitions
22662290
perform as purported.
22672291
@@ -2345,6 +2369,15 @@ def test_permutations_sizeof(self):
23452369
>>> first_true('ABC0DEF1', '9', str.isdigit)
23462370
'0'
23472371
2372+
>>> population = 'ABCDEFGH'
2373+
>>> for r in range(len(population) + 1):
2374+
... seq = list(combinations(population, r))
2375+
... for i in range(len(seq)):
2376+
... assert nth_combination(population, r, i) == seq[i]
2377+
... for i in range(-len(seq), 0):
2378+
... assert nth_combination(population, r, i) == seq[i]
2379+
2380+
23482381
"""
23492382

23502383
__test__ = {'libreftest' : libreftest}

0 commit comments

Comments
 (0)