嵌套字典列表的笛卡尔积

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}}

对于所有组合，都是如此......我正在寻找一种优雅的方法来实现这一点。没有特定的性能限制，它们是小字典。

- filippo

---

这是您需要的最深嵌套级别，还是您正在寻找适用于任意嵌套级别的解决方案？ - jdehesa

@jdehesa 对于我的当前用例来说，这是最深层嵌套的。 - filippo

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}}

- Martijn Pieters

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}}

- jdehesa

Answer 3

我扩展了Martijn Pieters的答案，使得函数还可以处理嵌套字典中的非可迭代值。因此下面的代码适用于类似这样的嵌套字典（只有一层嵌套）:

dictionary = {
    'a': 31415, 
    'b': 'strings too', 
    'c': [1, 2, 3], 
    'd': {'e': [1, 2, 3], 'f': 42, 'g': 'strings here too'}
}

from itertools import product

def _gen_combinations(d):
    keys, values = d.keys(), d.values()

    list_keys = [k for k in keys if isinstance(d[k], list)]
    nonlist_keys = [k for k in keys if k not in list_keys]
    list_values = [v for v in values if isinstance(v, list)]
    nonlist_values = [v for v in values if v not in list_values]

    combinations = product(*list_values)

    for c in combinations:
        result = dict(zip(list_keys, c))
        result.update({k: v for k, v in zip(nonlist_keys, nonlist_values)})
        yield result


def generate_dict_combinations(d):
    keys, values = d.keys(), d.values()

    dict_values = [v for v in values if isinstance(v, dict)]
    dict_keys = [k for k in keys if isinstance(d[k], dict)]
    nondict_values = [v for v in values if v not in dict_values]
    nondict_keys = [k for k in keys if k not in dict_keys]

    for c in product(*(_gen_combinations(v) for v in dict_values)):
        result = dict(zip(dict_keys, c))
        result.update({k: v for k, v in zip(nondict_keys, nondict_values)})
        yield result

- David Bellamy