kcombinations.py

'''calculate lists of indexes for n objects taken k at a time without respect for order'''

def indexlist1(n,k):
    '''
    soultion by counting up in binary with k pegs
    '''
    l = range(k)
    yield list(l)
    #until the first peg is in its maximum position
    while l[0] < (n - k):
        #identify first (from the left) unblocked (open space to the right) peg
        for peg in xrange(k):
            #final peg checks bounds of n
            #other pegs check the next peg
            if (peg == (k - 1) and l[peg] < (n - 1)) or (peg < (k - 1) and l[peg] < (l[peg+1] - 1)):
                #increment chosen peg
                l[peg] += 1
                #reset all lower indexed pegs to lowest position
                for lower_peg in range(peg):
                    l[lower_peg] = lower_peg
                break
        #yield a copy of the list 
        # so that modifications outside of this scope 
        # don't affect the iteration
        yield list(l)
    
def indexlist2(n,k):
    '''
    solution by counting to 2^n 
    and filtering out numbers with a binary 
    representation that containts exactly three ones
    '''
    for i in xrange(2**n):
        l = []
        length = 0
        shifts = 0
        #record indexes 
        # short circuit if there are more than k indexes
        # i is zero when there are no more indexes
        while i > 0 and length <= k:
            if i & 1:
                l.append(shifts)
                length += 1
            i = i >> 1
            shifts += 1
        if length == k:
            yield l

if __name__ == '__main__':
    solutions = [indexlist1,indexlist2]
    for s in solutions:
        print s.__name__
        print s.__doc__
        for x in s(5,3):
            print x
        print