python – Get the cartesian product of a series of lists?

python – Get the cartesian product of a series of lists?

itertools.product

Available from Python 2.6.

import itertools

somelists = [
   [1, 2, 3],
   [a, b],
   [4, 5]
]
for element in itertools.product(*somelists):
    print(element)

Which is the same as,

for element in itertools.product([1, 2, 3], [a, b], [4, 5]):
    print(element)
import itertools
>>> for i in itertools.product([1,2,3],[a,b],[4,5]):
...         print i
...
(1, a, 4)
(1, a, 5)
(1, b, 4)
(1, b, 5)
(2, a, 4)
(2, a, 5)
(2, b, 4)
(2, b, 5)
(3, a, 4)
(3, a, 5)
(3, b, 4)
(3, b, 5)
>>>

python – Get the cartesian product of a series of lists?

For Python 2.5 and older:

>>> [(a, b, c) for a in [1,2,3] for b in [a,b] for c in [4,5]]
[(1, a, 4), (1, a, 5), (1, b, 4), (1, b, 5), (2, a, 4), 
 (2, a, 5), (2, b, 4), (2, b, 5), (3, a, 4), (3, a, 5), 
 (3, b, 4), (3, b, 5)]

Heres a recursive version of product() (just an illustration):

def product(*args):
    if not args:
        return iter(((),)) # yield tuple()
    return (items + (item,) 
            for items in product(*args[:-1]) for item in args[-1])

Example:

>>> list(product([1,2,3], [a,b], [4,5])) 
[(1, a, 4), (1, a, 5), (1, b, 4), (1, b, 5), (2, a, 4), 
 (2, a, 5), (2, b, 4), (2, b, 5), (3, a, 4), (3, a, 5), 
 (3, b, 4), (3, b, 5)]
>>> list(product([1,2,3]))
[(1,), (2,), (3,)]
>>> list(product([]))
[]
>>> list(product())
[()]

Leave a Reply

Your email address will not be published.