I was asked to help implement a kind of reusable combination generator in Python given the maximum number and the group size. Combination is widely used in high performance computing especially for searching in state space. Note that it means solving problem by brute force technique. So I wrote several codes.

Recursive return list

This is a very simple version implemented using recursion.

def comblist(n,m,b=1,l=[0]):
    L = len(l)
    if L == m+1:
        return [l[1:]]
    else:
        c = l[-1]
        q,r = c+1,min(c+b+m,n+1)
        cl = []
        for i in range(q,r):
            cl.extend(comblist(n,m,b+1,l+[i]))
        return cl
 
print comblist(5,3)

Recursive with callback

Above code is good as long as the amount of combinations is not large since the whole list is stored in memory. I modified above code to use callback technique.

def comb(n,m,cb,b=1,l=[0]):
    L = len(l)
    if L == m+1:
        cb(l[1:])
    else:
        c = l[-1]
        q,r = c+1,min(c+b+m,n+1)
        for i in range(q,r):
            comb(n,m,cb,b+1,l+[i])

As a result, you can write separated callback for counting, collecting and etc.

class Counter:
    def __init__(self,value=0):
        self.value = value
 
    def count(self,v=1):
        self.value += v
 
cb = Counter()
comb(5,3,lambda l: cb.count())
print cb.value
 

class Collector(list):
    def __init__(self,a,b):
        super(Collector,self).__init__()
        self.first,self.last = a,b
        self.count = 0
 
    def add(self,l):
        if self.count >= self.first and self.count <= self.last:
            self.append(l)
        self.count += 1
 
cb = Collector(3,7)
comb(5,3,cb.add)
print cb

Self-contained Class

Anyway, you might want to obtain a list with a choice to extend using object-oriented technique. That's why I wrote below code.

class Combination(list):
    def __init__(self,n,m,a,b):
        super(Combination,self).__init__()
        self._n,self._m = n,m
        self._first,self._last = a,b
        self._count = 0
        self._generate()
 
    def _add(self,l):
        if self._count >= self._first and self._count <= self._last:
            self.append(l)
        self._count += 1
 
    def _generate(self,b=1,l=[0]):
        n,m = self._n,self._m
        L = len(l)
        if L == m+1:
            self._add(l[1:])
        else:
            c = l[-1]
            q,r = c+1,min(c+b+m,n+1)
            for i in range(q,r):
                self._generate(b+1,l+[i])
 
c = Combination(5,3,3,7)
print c

Tags: python, combination, codenone