Python itertools - 只创建所有可能产品的一个子集

Question

我正在编写一个脚本，该脚本使用 itertools 生成给定参数的所需组合。但是当一些变量被链接时，我在生成适当的“集合”时遇到了麻烦，即排除某些组合。请考虑以下事项:

import itertools

A = ['a1','a2','a3']
B = ['b1','b2','b3']
C = ['c1','c2']

如果我想生成这些元素的所有可能组合，我可以简单地使用 itertools.product()

all_combinations = list(itertools.product(A,B,C))

这给出了预期

[('a1', 'b1', 'c1'), ('a1', 'b1', 'c2'), ('a1', 'b2', 'c1'), ... 
('a3', 'b2', 'c2'), ('a3', 'b3', 'c1'), ('a3', 'b3', 'c2')]

对于 18 个组合 (3*3*2)

但是，我如何“链接”参数 A 和 B，使每个返回的集合仅包含“an”、“bn”元素？也就是说，我已经尝试过:

ABprecombine = zip(A,B)
limited_combinations = list(itertools.product(ABprecombine,C))

哪个返回

[(('a1', 'b1'), 'c1'), (('a1', 'b1'), 'c2'), (('a2', 'b2'), 'c1'),
 (('a2', 'b2'), 'c2'), (('a3', 'b3'), 'c1'), (('a3', 'b3'), 'c2')]

这是六个 (3*1*2) 所需的产品，但显然由于我创建它的方式，我现在有一个额外的元组。

当然，我可以生成所有组合，然后过滤掉给定的组合，但是是否有像上面那样“链接”参数的智能方法？

Answer 1

在这里，压缩 A 和 B 是正确的做法。如果需要，您可以很容易地展平元组:

limited_combinations = [(a, b, c) for ((a, b), c) in itertools.product(zip(A, B), C)]

如果您想更详细地控制产生哪些组合，事情会迅速变得更加复杂，直到需要解决 bool 可满足性等 NP-hard 问题的难度。如果发生这种情况，请查看现有库中的此类内容。