Python itertools.product 挑战：扩展包含元组的字典

Question

假设你有一个字典，里面有一些项目是元组（也就是一组数据）。

params = {
 'a': 'static',
 'b': (1, 2),
 'c': ('X', 'Y')
}

我需要把这些项目的“乘积”变成一个字典的列表，像这样，并且把元组展开，这样b里的每个项目都会和c里的每个项目匹配...

[{ 'a': 'static', 'b': 1, 'c': 'X' },
 { 'a': 'static', 'b': 1, 'c': 'Y' },
 { 'a': 'static', 'b': 2, 'c': 'X' },
 { 'a': 'static', 'b': 2, 'c': 'Y')}]

我可以很容易地把最开始的输入分成两部分，一部分是普通的项目，另一部分是元组项目。然后我会把每个元组的键当作“标签”应用到它的值上，这样它们看起来像这样：'b##1', 'b##2', 'c##X', 'c##Y'。接着在乘法运算后再把这些解析回上面的字典。如果我总是能看到两个元组项目（比如b和c），我可以很简单地把它们传给itertools.products。但是可能会有0到n个元组项目，而product()并不能以这种方式对一个列表的列表进行乘法运算。有没有人能想到解决办法？

TAG = '##'      
# separate tuples and non-tuples from the input, and prepend the key of each tuple as a tag on the value to parse out later
for key, value in params.items():
    if type(value) is tuple:
        for x in value:
            tuples.append(f'{key}{TAG}{x}')
    else:
        non_tuples.append({key: value})
print(list(product(tuples))      # BUG: doesn't distribute each value of b with each value of c

Answer 1

这里已经有人分享了明显的 itertools.product 解决方案，所以我来介绍一个快速的替代方法。这个方法一次展开一个参数（在线尝试这个！）:

params = {
    'a': 'static',
    'b': (1, 2),
    'c': ('X', 'Y')
}

out = [{}]
for k, v in params.items():
    if not isinstance(v, tuple):
        v = v,
    out = [d.copy()
           for d in out
           for d[k] in v]

print(out)

以下是三种情况的时间记录:

19 parameters with 2 options each:
  0.75 seconds  no_comment
  1.26 seconds  Andrej

10 parameters with 4 options each:
  0.60 seconds  no_comment
  1.63 seconds  Andrej

6 parameters with 10 options each:
  0.46 seconds  no_comment
  1.18 seconds  Andrej

基准测试脚本:

from itertools import product
from time import time
from string import ascii_lowercase

params = {c: (i, -i) for i, c in enumerate(ascii_lowercase[:19])}
params = {c: tuple(range(4)) for i, c in enumerate(ascii_lowercase[:10])}

def no_comment(params):
    out = [{}]
    for k, v in params.items():
        if not isinstance(v, tuple):
            v = v,
        out = [d.copy()
               for d in out
               for d[k] in v]
    return out

def Andrej(params):
    return [
        dict(zip(params, vals))
        for vals in product(
            *(v if isinstance(v, (tuple, list)) else (v,) for v in params.values())
        )
    ]

print(no_comment(params) == Andrej(params))

for f in [no_comment, Andrej] * 3:
    t0 = time()
    f(params)
    print(time() - t0, f.__name__)

在线尝试这个！

一个可能的优化隐藏在一个 markdown 注释中，查看答案的源代码可以看到它。

Answer 2

试试这个：

from itertools import product

params = {"a": "static", "b": (1, 2), "c": ("X", "Y")}


out = []
for vals in product(
    *[v if isinstance(v, (tuple, list)) else (v,) for v in params.values()]
):
    out.append(dict(zip(params, vals)))

print(out)

输出结果是：

[
    {"a": "static", "b": 1, "c": "X"},
    {"a": "static", "b": 1, "c": "Y"},
    {"a": "static", "b": 2, "c": "X"},
    {"a": "static", "b": 2, "c": "Y"},
]

---

一句话解决方案：

out = [
    dict(zip(params, vals))
    for vals in product(
        *(v if isinstance(v, (tuple, list)) else (v,) for v in params.values())
    )
]

Answer 3

product 可以接收多个可迭代对象，但要记住一个重要的点：可迭代对象可以只包含一个元素。如果你原来的字典中的某个值不是元组（或者可能是列表），你需要把它转换成一个只包含一个值的元组，然后再传给 product：

params_iterables = {}
for k, v in params.items():
    if isinstance(v, (tuple, list)):
        params_iterables[k] = v     # v is already a tuple or a list
    else:
        params_iterables[k] = (v, ) # A tuple containing a single value, v

这样就会得到：

params_iterables = {'a': ('static',), 'b': (1, 2), 'c': ('X', 'Y')}

接下来，只需计算 params_iterables 中值的乘积：

result = []
for values in product(*params_iterables.values()):
    result.append(dict(zip(params, values)))

dict(zip(params, values)) 这一行代码创建了一个字典，其中 values 的第一个元素被分配给 params 的第一个键，以此类推。这个字典随后被添加到 result 中，从而得到你想要的输出：

[{'a': 'static', 'b': 1, 'c': 'X'},
 {'a': 'static', 'b': 1, 'c': 'Y'},
 {'a': 'static', 'b': 2, 'c': 'X'},
 {'a': 'static', 'b': 2, 'c': 'Y'}]