Commit d37258dd authored by Raymond Hettinger's avatar Raymond Hettinger Committed by GitHub

Add itertools recipe for directly finding the n-th combination (#5161)

parent 0f31c74f
...@@ -859,6 +859,29 @@ which incur interpreter overhead. ...@@ -859,6 +859,29 @@ which incur interpreter overhead.
indices = sorted(random.randrange(n) for i in range(r)) indices = sorted(random.randrange(n) for i in range(r))
return tuple(pool[i] for i in indices) 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 Note, many of the above recipes can be optimized by replacing global lookups
with local variables defined as default values. For example, the with local variables defined as default values. For example, the
*dotproduct* recipe can be written as:: *dotproduct* recipe can be written as::
......
...@@ -2262,6 +2262,30 @@ Samuele ...@@ -2262,6 +2262,30 @@ Samuele
... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x ... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
... return next(filter(pred, iterable), default) ... 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 This is not part of the examples but it tests to make sure the definitions
perform as purported. perform as purported.
...@@ -2345,6 +2369,15 @@ True ...@@ -2345,6 +2369,15 @@ True
>>> first_true('ABC0DEF1', '9', str.isdigit) >>> first_true('ABC0DEF1', '9', str.isdigit)
'0' '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} __test__ = {'libreftest' : libreftest}
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment