辞書型データのvalueの直積を返したい

やりたいこと

pythonで、辞書が持つvalue同士の直積を、対応するkeyを変えずに辞書としてそれぞれ出力したい。

具体的には

{'A':[1,2,3], 'B':[4,5]}

という辞書が存在するときに

[{'A':1, 'B':4}, {'A':1, 'B':5},
 {'A':2, 'B':4}, {'A':2, 'B':5},
 {'A':3, 'B':4}, {'A':3, 'B':5}]

といった具合のリストが得られるようにしたい。

前提知識

itertools.product()で直積を求められます。

https://docs.python.org/ja/3/library/itertools.html#itertools.product

実装

test1.py

import itertools

test1 = {'A': [1, 2, 3], 'B': [4, 5]}

product = [x for x in itertools.product(*test1.values())]
result = [dict(zip(test1.keys(), r)) for r in product]
for r in result:
    print(r)

console

{'A': 1, 'B': 4}
{'A': 1, 'B': 5}
{'A': 2, 'B': 4}
{'A': 2, 'B': 5}
{'A': 3, 'B': 4}
{'A': 3, 'B': 5}

valueの型が異なっていてもOK

test2.py

import itertools

test2 = {'A': ['TEST', 1, 2.5], 'B': [[3, 4], 5]}

product = [x for x in itertools.product(*test2.values())]
result = [dict(zip(test1.keys(), r)) for r in product]
for r in result:
    print(r)

console

{'A': 'TEST', 'B': [3, 4]}
{'A': 'TEST', 'B': 5}
{'A': 1, 'B': [3, 4]}
{'A': 1, 'B': 5}
{'A': 2.5, 'B': [3, 4]}
{'A': 2.5, 'B': 5}

注意点

直積を求めたい辞書のvalueはlist型でなければなりません。
例えば以下の様な辞書を変換しようとするとエラーとなります。

test3.py

import itertools

test3 = {'A': 1, 'B': [2, 3], 'C': [4, 5, 6]}

product = [x for x in itertools.product(*test3.values())]
result = [dict(zip(test1.keys(), r)) for r in product]
for r in result:
    print(r)

console

Traceback (most recent call last):
  File "test3.py", line 5, in <module>
    product = [x for x in itertools.product(*test1.values())]
TypeError: 'int' object is not iterable

この様な場合は一旦valueを舐めてlistでなければlistに変換してあげれば良いです。

test4.py

import itertools

test4 = {'A': 1, 'B': [2, 3], 'C': [4, 5, 6]}

test4_after = dict([(key, val if type(val) is list else [val]) for key, val in test4.items()])

product = [x for x in itertools.product(*test4_after.values())]
result = [dict(zip(test4_after.keys(), r)) for r in product]
for r in result:
    print(r)

console

{'A': 1, 'B': 2, 'C': 4}
{'A': 1, 'B': 2, 'C': 5}
{'A': 1, 'B': 2, 'C': 6}
{'A': 1, 'B': 3, 'C': 4}
{'A': 1, 'B': 3, 'C': 5}
{'A': 1, 'B': 3, 'C': 6}

key値を捨てる場合

key値を捨ててvalueの直積だけ欲しい場合はzipの部分を消せばOK。

test5.py

import itertools

test5 = {'A': [1, 2], 'B': [4, 5, 6]}

product = [x for x in itertools.product(*test5.values())]
for r in product:
    print(r)

console

(1, 4)
(1, 5)
(1, 6)
(2, 4)
(2, 5)
(2, 6)