Search code examples
mathstatisticsprobabilitycombinatorics

Nth Combination


Is there a direct way of getting the Nth combination of an ordered set of all combinations of nCr?

Example: I have four elements: [6, 4, 2, 1]. All the possible combinations by taking three at a time would be: [[6, 4, 2], [6, 4, 1], [6, 2, 1], [4, 2, 1]].

Is there an algorithm that would give me e.g. the 3rd answer, [6, 2, 1], in the ordered result set, without enumerating all the previous answers?


Solution

  • Note you can generate the sequence by recursively generating all combinations with the first element, then all combinations without. In both recursive cases, you drop the first element to get all combinations from n-1 elements. In Python:

    def combination(l, r):
        if r == 0:
            yield []
        elif len(l) == r:
            yield l
        else:
            for c in (combination(l[1:], r-1)):
                yield l[0:1]+c
            for c in (combination(l[1:], r)):
                yield c
    

    Any time you're generating a sequence by making a choice like this, you can recursively generate the kth element by counting how many elements a choice generates and comparing the count to k. If k is less than the count, you make that choice. Otherwise, subtract the count and repeat for the other possible choices you could make at that point. If there are always b choices, you can view this as generating a number in base b. The technique still works if the number of choices varies. In pseudocode (when all choices are always available):

    kth(k, choicePoints)
        if choicePoints is empty
            return empty list
        for each choice in head of choicePoints:
            if k < size of choice
                return choice and kth(k, tail of choicePoints)
            else
                k -= size of choice
        signal exception: k is out-of-bounds
    

    This gives you a 0-based index. If you want 1-based, change the comparison to k <= size of choice.

    The tricky part (and what is unspecified in the pseudocode) is that the size of a choice depends on previous choices. Note the pseudocode can be used to solve a more general case than the problem.

    For this specific problem, there are two choices (b= 2) and the size of the 1st choice (i.e. including the 1st element) is given by n-1Cr-1. Here's one implementation (which requires a suitable nCr):

    def kthCombination(k, l, r):
        if r == 0:
            return []
        elif len(l) == r:
            return l
        else:
            i=nCr(len(l)-1, r-1)
            if k < i:
                return l[0:1] + kthCombination(k, l[1:], r-1)
            else:
                return kthCombination(k-i, l[1:], r)
    

    If you reverse the order of the choices, you reverse the order of the sequence.

    def reverseKthCombination(k, l, r):
        if r == 0:
            return []
        elif len(l) == r:
            return l
        else:
            i=nCr(len(l)-1, r)
            if k < i:
                return reverseKthCombination(k, l[1:], r)
            else:
                return l[0:1] + reverseKthCombination(k-i, l[1:], r-1)
    

    Putting it to use:

    >>> l = [6, 4, 2, 1]
    >>> [kthCombination(k, [6, 4, 2, 1], 3) for k in range(nCr(len(l), 3)) ]
    [[6, 4, 2], [6, 4, 1], [6, 2, 1], [4, 2, 1]]
    >>> powOf2s=[2**i for i in range(4,-1,-1)]
    >>> [sum(kthCombination(k, powOf2s, 3)) for k in range(nCr(len(powOf2s), 3))]
    [28, 26, 25, 22, 21, 19, 14, 13, 11, 7]
    >>> [sum(reverseKthCombination(k, powOf2s, 3)) for k in range(nCr(len(powOf2s), 3))]
    [7, 11, 13, 14, 19, 21, 22, 25, 26, 28]