Invertible combination generator for ordered lists

Question

I have code where given a list of ordered integers I generate several new lists, and I need to be able to map these lists one-to-one with the integers from 1 to N choose k (or 0 to (N choose k)-1) in python).

So for example if I have a N=7, k=3, then I might start with the list [0,1,2] which I might enumerate by 0. Then I generate the lists [0,1,3], [0,1,5], and [3,4,6] which all need to be uniquely identified by an integer in the range 0,...,34.

I know that I could use itertools to generate the ordered lists,

itertools.combinations(range(7), 3)

but I need an inverse function which is more efficient than simply searching through a pregenerated list since I will be working with large N and each list references around 50 new lists.

Answer 1

You can use this recipe from the itertools recipe here:

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)

Example usage:

>>> nth_combination(range(7), 3, 5)
(0, 2, 3)
>>> nth_combination(range(7), 3, 34)
(4, 5, 6)

To reverse:

from math import factorial
def get_comb_index(comb, n):
    k = len(comb)
    rv = 0
    curr_item = 0
    n -= 1
    for offset, item in enumerate(comb, 1):
        while curr_item < item:
            rv += factorial(n-curr_item)//factorial(k-offset)//factorial(n+offset-curr_item-k)
            curr_item += 1
        curr_item += 1
    return rv

Example Usage:

>>> get_comb_index((4,5,6), 7)
34
>>> get_comb_index((0,1,2), 7)
0
>>> get_comb_index((0,2,4), 7)
6