我需要生成n组的所有k子集,并附加约束,即必须一起选择某些元素对或根本不选择它们。为了模拟这个约束,我想到明确地将这些元素配对为2元组并将其他元素保持为1元组。 因此,例如,假设我需要选择{1,2,3,4,5}的所有3元素子集,并且必须一起选择元素3和4的附加约束。然后我的新集是:
{(1,), (2,), (3, 4), (5,)}
我要编写的函数需要生成:
{1, 2, 5}, {1, 3, 4}, {2, 3, 4}, {3, 4, 5}.
有没有一种简单的方法可以使用itertools(或者可能不知道的其他python模块)来获得这个结果?我不关心接收这些子集的顺序。
如果这简化了事情:一个元素不能与多个其他元素配对(所以(3,5)例如在我的例子中不能作为附加约束出现。)
答案 0 :(得分:1)
解决方案:
from itertools import combinations, chain
def faster(pairs, others, k):
for npairs in range(k // 2 + 1):
for pairs_comb in combinations(pairs, npairs):
for others_comb in combinations(others, k - npairs * 2):
yield chain(others_comb, *pairs_comb)
说明:
完成结果中对数的所有可能性。例如,如果k = 5
那么可以没有对和5个无约束元素(others
),或者1对和3个其他元素,或者2对和1个其他元素。然后,所有组合和其他组合可以独立生成并组合。
测试:
def brute_force(pairs, others, k):
return [c for c in combinations(chain(others, *pairs), k)
if all((p1 in c) == (p2 in c) for p1, p2 in pairs)]
def normalise(combs):
return sorted(map(sorted, combs))
args = ([(3, 4), (1, 2), (6, 7)], [5, 8, 9, 10, 11], 4)
assert normalise(brute_force(*args)) == normalise(faster(*args))
print(normalise(faster(*args)))