生成具有配对约束的所有子集

Question

我需要生成一个 n-set 的所有 k-subsets，附加约束是必须一起选择一些元素对，或者根本不选择。为了对该约束建模，我考虑将这些元素显式配对为 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) 在我的示例中不能作为附加约束出现）。

Answer 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 个其他元素。那么所有pair和other的组合都可以独立生成和组合。

测试：

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