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

rhettinger · web-flow · commit d37258dd2e18 · 2018-01-13T10:35:40.000-08:00
diff --git 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
@@ -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}
