comb.py

def comb(array, k):
    """
    Get all combinations of k elements from array

    Arguments: 
    array -- an array of elements, source for combinations
    k -- number of elements in combination

    Return:
    comb -- array of stored combinations

    Examples:
    comb([1, 2, 3], 2) 
        returns [[1, 2], [1, 3], [2, 3]]

    comb([1, 2, 3, 4], 3)
        returns [[1, 2, 3], [1, 2, 4], [2, 3, 4]]
    
    """
    def comb_internal(array, k, level, start_indx, path):
        result = []
        if level == k:
            return [path]

        for i in range(start_indx, 1 + len(array) - k + level):
            temp = comb_internal(array, k, level + 1, i + 1, path + [array[i]])
            result += temp

        return result

    return comb_internal(array, k, 0, 0, [])

## tests

import numpy as np
import math

# check if number of combination is equal to n! / (k! * (n-k)!) 
assert(
    ( lambda array, k:
        len(comb(array, k))
            == math.factorial(len(array)) / math.factorial(k) / math.factorial(len(array) - k)
    ) ([1, 2, 3, 4, 5], 3)
)

# check if number of unique sorted combinations is the same as total number of combinations
assert(
    ( lambda array, k:
        len(comb(array, k))
            == len(
                np.unique(
                    list(
                       map(
                           lambda x: sorted(x), 
                           comb(array, k)
                       )
                    ),
                    axis = 0
                )
            )
    ) ([1, 2, 3, 4, 5, 6, 7, 8], 4)
)

# check bounding conditions
assert(comb([1, 2, 3, 4], 1) == [[1], [2], [3], [4]])

## TODO: fix this assert, currently comb([], 0) returns [[]], instead of []
# assert(comb([], 0) == [])

assert(comb([], 2) == [])

assert(comb(['a', 'b', 'c'], 2) == [['a', 'b'], ['a', 'c'], ['b', 'c']])

assert(comb([1, 2], 10) == [])

assert(comb([1, 2, 3], 3) == [[1, 2, 3]])