列表嵌套字典的笛卡尔积

Question

我有一些代码可以生成列表字典的所有组合

import itertools
import collections

def gen_combinations(d):
    keys, values = d.keys(), d.values()
    combinations = itertools.product(*values)

    for c in combinations:
        yield dict(zip(keys, c))

假设我有一本字典 A 就像

A = {'a': [0, 1],
     'b': [2, 3, 4]}

它产生：

{'a': 0, 'b': 2}
{'a': 0, 'b': 3}
{'a': 0, 'b': 4}
{'a': 1, 'b': 2}
{'a': 1, 'b': 3}
{'a': 1, 'b': 4}

现在我希望它能与列表的嵌套字典一起工作，例如

B = {'s1': {'a': [0, 1],
            'b': [0, 1, 2] },
     's2': {'c': [0, 1],
            'd': [0, 1] }}

应该会产生如下内容：

{'s1': {'a': 0, 'b': 0},
 's2': {'c': 0, 'd': 0}}

{'s1': {'a': 0, 'b': 0},
 's2': {'c': 0, 'd': 1}}

{'s1': {'a': 0, 'b': 0},
 's2': {'c': 1, 'd': 0}}

{'s1': {'a': 0, 'b': 0},
 's2': {'c': 1, 'd': 1}}

等等所有的组合... 我正在寻找一种优雅的方式来做到这一点。没有任何特定的性能限制，它们是小型词典。

Answer 1

只需为外部字典中的每个键创建 gen_combinations() 输出的乘积：

def gen_dict_combinations(d):
    keys, values = d.keys(), d.values()
    for c in itertools.product(*(gen_combinations(v) for v in values)):
        yield dict(zip(keys, c))

这基本上是相同的模式，只是现在我们不直接使用 values，而是使用 gen_combinations():

的输出

>>> for c in gen_dict_combinations(B):
...     print(c)
...
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 1, 'd': 1}}

Answer 2

最终得到了任意嵌套字典的递归版本：

import itertools

def gen_combinations(d):
    keys, values = d.keys(), d.values()
    values_choices = (gen_combinations(v) if isinstance(v, dict) else v for v in values)
    for comb in itertools.product(*values_choices):
        yield dict(zip(keys, comb))


B = {'s1': {'a': [0, 1],
            'b': [0, 1, 2] },
     's2': {'c': [0, 1],
            'd': [0, 1] }}

for c in gen_combinations(B):
     print(c)

输出：

{'s1': {'a': 0, 'b': 0}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 0, 'b': 0}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 0, 'b': 1}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 0, 'b': 2}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 1, 'b': 0}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 1, 'b': 1}, 's2': {'c': 1, 'd': 1}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 0, 'd': 0}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 0, 'd': 1}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 1, 'd': 0}}
{'s1': {'a': 1, 'b': 2}, 's2': {'c': 1, 'd': 1}}

列表嵌套字典的笛卡尔积

Cartesian product of nested dictionaries of lists

python

combinations

cartesian-product

itertools

python-3.x