Add itertools recipe for directly finding the n-th combination (#5161)
This commit is contained in:
parent
0f31c74fcf
commit
d37258dd2e
|
@ -859,6 +859,29 @@ which incur interpreter overhead.
|
|||
indices = sorted(random.randrange(n) for i in range(r))
|
||||
return tuple(pool[i] for i in indices)
|
||||
|
||||
def nth_combination(iterable, r, index):
|
||||
'Equivalent to list(combinations(iterable, r))[index]'
|
||||
pool = tuple(iterable)
|
||||
n = len(pool)
|
||||
if r < 0 or r > n:
|
||||
raise ValueError
|
||||
c = 1
|
||||
k = min(r, n-r)
|
||||
for i in range(1, k+1):
|
||||
c = c * (n - k + i) // i
|
||||
if index < 0:
|
||||
index += c
|
||||
if index < 0 or index >= c:
|
||||
raise IndexError
|
||||
result = []
|
||||
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)
|
||||
|
||||
Note, many of the above recipes can be optimized by replacing global lookups
|
||||
with local variables defined as default values. For example, the
|
||||
*dotproduct* recipe can be written as::
|
||||
|
|
|
@ -2262,6 +2262,30 @@ Samuele
|
|||
... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
|
||||
... return next(filter(pred, iterable), default)
|
||||
|
||||
>>> def nth_combination(iterable, r, index):
|
||||
... 'Equivalent to list(combinations(iterable, r))[index]'
|
||||
... pool = tuple(iterable)
|
||||
... n = len(pool)
|
||||
... if r < 0 or r > n:
|
||||
... raise ValueError
|
||||
... c = 1
|
||||
... k = min(r, n-r)
|
||||
... for i in range(1, k+1):
|
||||
... c = c * (n - k + i) // i
|
||||
... if index < 0:
|
||||
... index += c
|
||||
... if index < 0 or index >= c:
|
||||
... raise IndexError
|
||||
... result = []
|
||||
... 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)
|
||||
|
||||
|
||||
This is not part of the examples but it tests to make sure the definitions
|
||||
perform as purported.
|
||||
|
||||
|
@ -2345,6 +2369,15 @@ True
|
|||
>>> first_true('ABC0DEF1', '9', str.isdigit)
|
||||
'0'
|
||||
|
||||
>>> population = 'ABCDEFGH'
|
||||
>>> for r in range(len(population) + 1):
|
||||
... seq = list(combinations(population, r))
|
||||
... for i in range(len(seq)):
|
||||
... assert nth_combination(population, r, i) == seq[i]
|
||||
... for i in range(-len(seq), 0):
|
||||
... assert nth_combination(population, r, i) == seq[i]
|
||||
|
||||
|
||||
"""
|
||||
|
||||
__test__ = {'libreftest' : libreftest}
|
||||
|
|
Loading…
Reference in New Issue