Cartesian product of a dictionary of lists
I'm trying to write some code to test out the Cartesian product of a bunch of input parameters.
I've looked at itertools
, but its product
function is not exactly what I want. Is there a simple obvious way to take a dictionary with an arbitrary number of keys and an arbitrary number of elements in each value, and then yield a dictionary with the next permutation?
Input:
options = {"number": [1,2,3], "color": ["orange","blue"] }
print list( my_product(options) )
Example output:
[ {"number": 1, "color": "orange"},
{"number": 1, "color": "blue"},
{"number": 2, "color": "orange"},
{"number": 2, "color": "blue"},
{"number": 3, "color": "orange"},
{"number": 3, "color": "blue"}
]
Ok, thanks @dfan for telling me I was looking in the wrong place. I've got it now:
from itertools import product
def my_product(inp):
return (dict(zip(inp.keys(), values)) for values in product(*inp.values())
EDIT: after years more Python experience, I think a better solution is to accept kwargs
rather than a dictionary of inputs; the call style is more analogous to that of the original itertools.product
. Also I think writing a generator function, rather than a function that returns a generator expression, makes the code clearer. So:
def product_dict(**kwargs):
keys = kwargs.keys()
vals = kwargs.values()
for instance in itertools.product(*vals):
yield dict(zip(keys, instance))
and if you need to pass in a dict, list(product_dict(**mydict))
. The one notable change using kwargs
rather than an arbitrary input class is that it prevents the keys/values from being ordered, at least until Python 3.6.
Python 3 version of Seth's answer.
import itertools
def dict_product(dicts):
"""
>>> list(dict_product(dict(number=[1,2], character='ab')))
[{'character': 'a', 'number': 1},
{'character': 'a', 'number': 2},
{'character': 'b', 'number': 1},
{'character': 'b', 'number': 2}]
"""
return (dict(zip(dicts, x)) for x in itertools.product(*dicts.values()))