[Python-checkins] Add itertools recipe for directly finding the n-th combination (#5161)

Sat Jan 13 13:35:42 EST 2018

https://github.com/python/cpython/commit/d37258dd2e189141906bd234385096cd8e885d8d
commit: d37258dd2e189141906bd234385096cd8e885d8d
branch: master
author: Raymond Hettinger <rhettinger at users.noreply.github.com>
committer: GitHub <noreply at github.com>
date: 2018年01月13日T10:35:40-08:00
summary:
Add itertools recipe for directly finding the n-th combination (#5161)
files:
M Doc/library/itertools.rst
M Lib/test/test_itertools.py

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 5af5422eefe..0b3829f19fa 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -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::
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index a359905e2e4..4fcc02acbf4 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -2262,6 +2262,30 @@ def test_permutations_sizeof(self):
 ... # 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 @@ def test_permutations_sizeof(self):
 >>> 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}
