Overview
The key to this is to walk through the cumulative combinations until the index is reached.
Solution
from math import factorial
def comb(n, r):
'Combinations of n things taken r at a time'
return factorial(n) // (factorial(r) * factorial(n - r))
def nth_combination(population, r, index):
'Equivalent to list(combinations(population, r))[index]'
n = len(population)
c = comb(n, r)
if index < 0:
index += c
if index < 0 or index >= c:
raise IndexError
if r < 0 or r > n:
raise ValueError
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(population[-1-n])
return tuple(result)
Optimization
If speed is a concern, it is possible to to build faster version of the comb() function.
One way is to precompute the factorials and then look them up when needed:
fact = [1]
for i in range(1, 1000):
fact.append(fact[-1] * i)
def comb(n, r):
return fact[n] // (fact[r] * fact[n - r])
And there is another way that avoids large factorials altogether and doesn't require auxiliary storage:
def comb(n, r):
c = 1
r = min(r, n-r)
for i in range(1, r+1):
c = c * (n - r + i) // i
return c
How it works
Start by breaking the combinations into its component groups:
def groups(n, r):
return [comb(n-i-1, r-1) for i in range(n-r+1)]
>>> comb(8, 3)
56
>>> groups(8, 3)
[21, 15, 10, 6, 3, 1]
This means that when you run itertools.combinations('ABCDEFGH', 3)
for n=8
letters taken r=3
at a time, there are 56 combinations. The first 21 start with A
, the next 15 start with B
, the next 10 start with C
, the next 6 start with D
, the next 3 start with E
, and the last 1 starts with F
.
Say you want to find the 25th combination out of the 56. That falls in the second group, so your first letter is B
.
And since 25 - 21 is 4, you'll then need to find the 4th combination in the 15 members of "B" group defined by itertools.combinations('CDEFGH', 2)
. Just repeat the above process until all the letters are extracted.
Testing
Here's a test to make sure it gives the expected results:
from itertools import combinations
population = 'ABCDEFGH'
for r in range(len(population) + 1):
for i, expected in enumerate(combinations(population, r)):
actual = locate(list(population), r, i)
assert expected == actual