Wednesday, March 30, 2011

Get the cartesian product of a series of lists in Python

How can I get the Cartesian product (every possible combination of values) from a group of lists?

Input:

somelists = [
   [1, 2, 3],
   ['a', 'b'],
   [4, 5]
]

Desired output:

[(1, 'a', 4), (1, 'a', 5), (1, 'b', 4), (1, 'b', 5), (2, 'a', 4), (2, 'a', 5) ...]
From stackoverflow
  • In Python 2.6+

    import itertools
    for element in itertools.product(*somelists):
        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)
    >>>
    
  • with itertools.product:

    import itertools
    result = list(itertools.product(*somelists))
    
  • In Python 2.6 and above you can use 'itertools.product`. In older versions of Python you can use the following (almost -- see documentation) equivalent code from the documentation:

    def product(*args, **kwds):
        # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
        # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
        pools = map(tuple, args) * kwds.get('repeat', 1)
        result = [[]]
        for pool in pools:
            result = [x+[y] for x in result for y in pool]
        for prod in result:
            yield tuple(prod)
    

    The result of both is an iterator, so if you really need a list for furthert processing, use list(result).

    Triptych : Per the documentation, the actual itertools.product implementation does NOT build intermediate results, which could be expensive. Using this technique could get out of hand quite quickly for moderately sized lists.
    hop : i can only point the OP to the documentation, not read it for him.
    hop : i forgot... programmers are babies
    recursive : Cute babies! Nevar forget.
    Triptych : The code from the documentation is meant to demonstrate what the product function does, not as a workaround for earlier versions of Python.
    hop : your point being?
  • 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)]
    

    Here's 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())
    [()]
    
    J.F. Sebastian : The recursive version doesn't work if some of `args` are iterators.

0 comments:

Post a Comment