Search code examples
pythonarraysdictionarycartesian-product

Cartesian product of dict of lists in Python


I would like to have a Python function cartesian_product which takes a dictionary of lists as input and returns as output a list containing as elements all possible dictionary which can be formed by taking from each list one element. Here would be an example:

Calling

cartesian_product({1: ['a', 'b'], 2: ['c', 'd'], 3: ['e', 'f']})

should return

[
  {1: 'a', 2: 'c', 3: 'e'},
  {1: 'a', 2: 'c', 3: 'f'},
  {1: 'a', 2: 'd', 3: 'e'},
  {1: 'a', 2: 'd', 3: 'f'},
  {1: 'b', 2: 'c', 3: 'e'},
  {1: 'b', 2: 'c', 3: 'f'},
  {1: 'b', 2: 'd', 3: 'e'},
  {1: 'b', 2: 'd', 3: 'f'}
]

Solution

  • This should do the trick:

    import itertools
    
    def cartesian_product(d):
        return [dict(zip(d, p)) for p in itertools.product(*d.values())]
    

    Demo:

    >>> d = {1: ['a', 'b'], 2: ['c', 'd'], 3: ['e', 'f']}
    >>> from pprint import pp
    >>> pp(cartesian_product(d))
    [{1: 'a', 2: 'c', 3: 'e'},
     {1: 'a', 2: 'c', 3: 'f'},
     {1: 'a', 2: 'd', 3: 'e'},
     {1: 'a', 2: 'd', 3: 'f'},
     {1: 'b', 2: 'c', 3: 'e'},
     {1: 'b', 2: 'c', 3: 'f'},
     {1: 'b', 2: 'd', 3: 'e'},
     {1: 'b', 2: 'd', 3: 'f'}]