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

Question

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

Answer 1

I dug up some old (although it's been converted to Python 3 syntax) code that includes the function combination_index which does what you request:

def fact(n, _f=[1, 1, 2, 6, 24, 120, 720]):
    """Return n!
    The “hidden” list _f acts as a cache"""
    try:
        return _f[n]
    except IndexError:
        while len(_f) <= n:
            _f.append(_f[-1] * len(_f))
        return _f[n]

def indexed_combination(n: int, k: int, index: int) -> tuple:
    """Select the 'index'th combination of k over n
    Result is a tuple (i | i∈{0…n-1}) of length k

    Note that if index ≥ binomial_coefficient(n,k)
    then the result is almost always invalid"""

    result= []
    for item, n in enumerate(range(n, -1, -1)):
        pivot= fact(n-1)//fact(k-1)//fact(n-k)
        if index < pivot:
            result.append(item)
            k-= 1
            if k <= 0: break
        else:
            index-= pivot
    return tuple(result)

def combination_index(combination: tuple, n: int) -> int:
    """Return the index of combination (length == k)

    The combination argument should be a sorted sequence (i | i∈{0…n-1})"""

    k= len(combination)
    index= 0
    item_in_check= 0
    n-= 1 # to simplify subsequent calculations
    for offset, item in enumerate(combination, 1):
        while item_in_check < item:
            index+= fact(n-item_in_check)//fact(k-offset)//fact(n+offset-item_in_check-k)
            item_in_check+= 1
        item_in_check+= 1
    return index

def test():
    for n in range(1, 11):
        for k in range(1, n+1):
            max_index= fact(n)//fact(k)//fact(n-k)
            for i in range(max_index):
                comb= indexed_combination(n, k, i)
                i2= combination_index(comb, n)
                if i2 != i:
                    raise RuntimeError("mismatching n:%d k:%d i:%d≠%d" % (n, k, i, i2))

indexed_combination does the inverse operation.

PS I remember that I sometime attempted removing all those fact calls (by substituting appropriate incremental multiplications and divisions) but the code became much more complicated and wasn't actually faster. A speedup was achievable if I substituted a pre-calculated list of factorials for the fact function, but again the speed difference was negligible for my use cases, so I kept this version.

Answer 2

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):
    ...

Answer 3

Looks like you need to better specify your task or I am just getting it wrong. For me it seems that when you iterating through the itertools.combination you can save indexes you need to an appropriate data structure. If you need all of them then I would go with the dict (one dict for all your needs):

combinationToIdx = {}
for (idx, comb) in enumerate(itertools.combinations(range(1,14),6)):
    combinationToIdx[comb] = idx

def findIdx(comb):
    return combinationToIdx[comb]