mmj mmj - 2 months ago 13
Python Question

Find the index of a given combination (of natural numbers) among those returned by `itertools` Python module

Given a combination of

k
of the first
n
natural numbers, for some reason I need to find the position of such combination among those returned by
itertools.combination(range(1,n),k)
(the reason is that this way I can use a
list
instead of a
dict
to access values associated to each combination, knowing the combination).

Since
itertools
yields its combinations in a regular pattern it is possible to do it (and I also found a neat algorithm), but I'm looking for an even faster/natural way which I might ignore.

By the way here is my solution:

def find_idx(comb,n):
k=len(comb)
idx=0
last_c=0
for c in comb:
#idx+=sum(nck(n-2-x,k-1) for x in range(c-last_c-1)) # a little faster without nck caching
idx+=nck(n-1,k)-nck(n-c+last_c,k) # more elegant (thanks to Ray), faster with nck caching
n-=c-last_c
k-=1
last_c=c
return idx


where
nck
returns the binomial coefficient of n,k.

For example:

comb=list(itertools.combinations(range(1,14),6))[654] #pick the 654th combination
find_idx(comb,14) # -> 654


And here is an equivalent but maybe less involved version (actually I derived the previous one from the following one). I considered the integers of the combination
c
as positions of 1s in a binary digit, I built a binary tree on parsing 0/1, and I found a regular pattern of index increments during parsing:

def find_idx(comb,n):
k=len(comb)
b=bin(sum(1<<(x-1) for x in comb))[2:]
idx=0
for s in b[::-1]:
if s=='0':
idx+=nck(n-2,k-1)
else:
k-=1
n-=1
return idx

Ray Ray
Answer

Your solution seems quite fast. In find_idx, you have two for loop, the inner loop can be optimized using the formular:

C(n, k) + C(n-1, k) + ... + C(n-r, k) = C(n+1, k+1) - C(n-r, k+1)

so, you can replace sum(nck(n-2-x,k-1) for x in range(c-last_c-1)) with nck(n-1, k) - nck(n-c+last_c, k).

I don't know how you implement your nck(n, k) function, but it should be O(k) measured in time complexity. Here I provide my implementation:

from operator import mul
from functools import reduce # In python 3
def nck_safe(n, k):
    if k < 0 or n < k: return 0
    return reduce(mul, range(n, n-k, -1), 1) // reduce(mul, range(1, k+1), 1)

Finally, your solution become O(k^2) without recursion. It's quite fast since k wouldn't be too large.

Update

I've noticed that nck's parameters are (n, k). Both n and k won't be too large. We may speed up the program by caching.

def nck(n, k, _cache={}):
    if (n, k) in _cache: return _cache[n, k]
    ....
    # before returning the result
    _cache[n, k] = result
    return result

In python3 this can be done by using functools.lru_cache decorator:

@functools.lru_cache(maxsize=500)
def nck(n, k):
    ...
Comments