Compute the nth combination in lexicographic order more efficiently, also with code for varying lengths

combination.py

#! /usr/bin/python
#
# This snippet provides a simple function "combination" that can compute an
# arbitrary k-combination of n elements given an index m into the
# lexicographically ordered set of all k-combinations of n elements.


from math import comb


# Compute the mth combination in lexicographical order from a set of n
# elements chosen k at a time.
# Algorithm from http://msdn.microsoft.com/en-us/library/aa289166(v=vs.71).aspx
def combination(n, k, m):
    result = []
    a = n
    b = k
    x = (comb(n, k) - 1) - m
    for i in range(0, k):
        a = a - 1
        while comb(a, b) > x:
            a = a - 1
        result.append(n - 1 - a)
        x = x - comb(a, b)
        b = b - 1
    return result


# Compute the mth combination in lexicographical order from a set of n elements of varying lengths
def unrankVaryingLengthCombination(n, rank):
    # Find length of k such that the rank is within the space of k-combination
    cumulative = 0
    found_k = 0
    for k in range(1, n + 1):
        cumulative += comb(n, k)
        found_k = k
        if rank < cumulative:
            break
    # Find rank within k-combinations from rank in 2^n - 1
    # Let's say C(4, 2) in 2^4 - 1 set, let's checkout rank of combination 4
    # Original combinoid is 4, should result in [0,1].
    # Cumulative would be 10 since C(4,1) + C(4,2) = 4+6 = 10
    # orignal_rank - cumulative = 4 - 10 = -6
    # We can see that resulting number is like inverted index for combinoid in C(4,2)
    # Proper rank should be in range of 0...C(4,2) which is 0...6
    # So we just add C(4,2) on top of our newly inverted rank: -6 + 6 = 0
    # Same would go for any consecutive number.
    # For example let's examine rank 5 in 2^4:
    # Original rank is 5, this should result in combination [0,2] in C(4,2) set
    # Cumulative is 10 since C(4,1) + C(4,2) = 4 + 6 = 10
    # 5 - 10 = -5 is inverted index of combinoid in C(4,2)
    # -5 + C(4,2) = -5 + 6 = 1 which is correct combinoid in C(4,2)
    rank = (rank - cumulative) + comb(n, found_k)

    # Just a debug check to make sure algorithm will not fail, should never be called
    if rank >= comb(n, found_k):
        raise ValueError(f"New Rank {rank} is greater than or equal to the maximum rank for C({n},{found_k}) = {comb(n, found_k)}")

    return combination(n, found_k, rank)


# Test the algorithm by printing all 3-combinations of 5 elements.  The
# algorithm is inefficient when done iteratively, it is designed primarily to
# find an arbitrary element in lexicographic order without having to iterate.
print("Test algorithm for mth combination in C(n,k)")
for i in range(0, comb(5, 3)):
    print(f"rank: {i} combination: {combination(5, 3, i)}")

# Test the algorithm by printing all combinations of varing lengths of 5 elements
print("Test algorithm for mth combination in 2^n - 1")
for i in range(0, 2 ** 5 - 1):
    print(f"rank: {i} combination: {unrankVaryingLengthCombination(5, i)}")