Combinations of n Lists

Remco Gerlich scarblac at pino.selwerd.nl
Thu Mar 1 00:10:45 CET 2001


Warren Postma <embed at geocities.com> wrote in comp.lang.python:
> I was just wondering if anyone has a more general version of this little
> helper function:
> 
> def combinations(list1,list2):
>     return [ (i,j) for i in list1 for j in list2 ]
> 
> print combinations( [1,2,3], ['a','b','c'] )
> 
> [(1, 'a'), (1, 'b'), (1, 'c'), (2, 'a'), (2, 'b'), (2, 'c'), (3, 'a'), (3,
> 'b'), (3, 'c')]
> 
> So, what if I want
> 
> def combinations(*lists):
>         ....
> print combinations( [1,2,3], [4,5,6], [7,8,9], .... )
>     [ [1,4,7,....],  .... [1,4,8, ... ], ..... ]
> 
> Any takers!? ;-)

Let's see, obvious tail recursive solution...

def combinations1(*lists):
   if len(lists) == 0:
      return []
   elif len(lists) == 1:
      return [(x,) for x in lists[0]]
   else: 
      return [(x,)+y for x in lists[0] for y in combinations(lists[1:])]

Change to iterative version:

def combinations2(*lists):
   if len(lists) == 0:
      return []
   result = [(x,) for x in lists[0]]
   for next in lists[1:]:
      result = [x+(y,) for x in result for y in next]
   return result
   
Well, that was obvious. Probably full of errors too, knowing myself (can't
test here). Hard to believe this is going to be the fastest way though,
building all those lists. Let's do something clever based on the fact that
every combination can be thought of as a number between 0 and
len(list0)*len(list1)*...*len(listn).

def combinations3(*lists):
   import operator
   lengths = map(len, lists)
   modulos = [reduce(operator.mul, lengths[:i],1) for i in range(len(lengths))]
   
   tuples = reduce(operator.mul, lengths)
   result = [None]*tuples
   for i in range(tuples):
      result[i] = tuple([lists[(i/modulos[list])%(lengths[list])] 
                           for list in range(len(lists))])
   return result

Hmm, this is a bit too sick, I doubt it's going to work at once, but I'm
curious what its speed is compared to the others :). Alas, can't test until
tomorrow.

If you have, say, six lists of lengths 6,5,7,8,9,10, then tuple nr. 5432 is
going to have item (5432/(6*5*7*8))%9 = 3 chosen from the fifth list, and so
on. 'modulos' is therefore misnamed, but it stores the number '6*5*7*8' for
the fifth list, etc.

Whatever, not going to write more docs now, time to sleep :).

I think this version is going to get trouble quickly when the amount of
tuples is more than maxint. But I think lists get trouble then anyway.

Thanks for the fun assignment :)

-- 
Remco Gerlich



More information about the Python-list mailing list