Merge e200855 into f450ab4

docs/api.rst

@@ -185,6 +185,7 @@ These tools yield combinatorial arrangements of items from iterables.
 .. autofunction:: random_permutation
 .. autofunction:: random_combination
 .. autofunction:: random_combination_with_replacement
+.. autofunction:: nth_combination
 
 
 Wrapping

more_itertools/recipes.py

@@ -28,6 +28,7 @@
     'iter_except',
     'ncycles',
     'nth',
+    'nth_combination',
     'padnone',
     'pairwise',
     'partition',
@@ -511,3 +512,39 @@ def random_combination_with_replacement(iterable, r):
     n = len(pool)
     indices = sorted(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]``.
+
+    The subsequences of *iterable* that are of length *r* can be ordered
+    lexicographically. :func:`nth_combination` computes the subsequence at
+    sort position *index* directly, without computing the previous
+    subsequences.
+
+    """
+    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)

more_itertools/tests/test_recipes.py

@@ -1,6 +1,7 @@
 from doctest import DocTestSuite
 from unittest import TestCase
 
+from itertools import combinations
 from six.moves import range
 
 import more_itertools as mi
@@ -574,3 +575,17 @@ def test_pseudorandomness(self):
             combination = mi.random_combination_with_replacement(items, 5)
             all_items |= set(combination)
         self.assertEqual(all_items, set(items))
+
+
+class NthCombinationTests(TestCase):
+    def test_basic(self):
+        iterable = 'abcdefg'
+        r = 4
+        for index, expected in enumerate(combinations(iterable, r)):
+            actual = mi.nth_combination(iterable, r, index)
+            self.assertEqual(actual, expected)
+
+    def test_long(self):
+        actual = mi.nth_combination(range(180), 4, 2000000)
+        expected = (2, 12, 35, 126)
+        self.assertEqual(actual, expected)