python - 在Python中生成成功集的所有组合

Question

有 k 个处理和 N 个总测试分布在这些处理中，这称为一个计划。对于一个固定的计划，我想用 Python 输出所有可能的成功集。

问题:

例如，如果医生正在测试头痛药，如果 k=2 种治疗方法(即阿司匹林和布洛芬)并且 N=3 总测试，一个计划可能是(1 次阿司匹林测试，2 次布洛芬测试)。对于该计划，我如何输出 0-1 次成功的阿司匹林测试和 0-2 次成功的布洛芬测试的所有可能组合？一项成功的测试意味着当给头痛患者服用阿司匹林时，阿司匹林可以治愈他们的头痛。

请用 python 代码发布答案，而不是数学答案。

期望的输出是一个列表，其中包含 [# successes for treatment 1, # successes of treatment 2]:

[ [0,0], [0,1], [0,2], [1,0], [1,1], [1,2] ]

如果可以使用 yield 就好了，因为上面的列表可能很长，我不想将整个列表存储在内存中，这会增加计算时间。

下面我有枚举 A 盒子中 N 个球的所有可能组合的代码，我认为这应该类似于创建所有可能的成功集，但我不确定如何实现。

代码

#Return list of tuples of all possible plans (n1,..,nk), where N = total # of tests = balls, K = # of treatments = boxes
#Code: Glyph, http://stackoverflow.com/questions/996004/enumeration-of-combinations-of-n-balls-in-a-boxes
def ballsAndBoxes(balls, boxes, boxIndex=0, sumThusFar=0):
    if boxIndex < (boxes - 1):
        for counter in range(balls + 1 - sumThusFar):
            for rest in ballsAndBoxes(balls, boxes,
                                      boxIndex + 1,
                                      sumThusFar + counter):
                yield (counter,) + rest
    else:
        yield (balls - sumThusFar,)

Answer 1

生成计划是一个划分问题，但是为给定计划生成成功集只需要生成一组范围的笛卡尔积。

from itertools import product

def success_sets(plan):
    return product(*map(lambda n: range(n + 1), plan))

plan = [1, 2]
for s in success_sets(plan):
    print(s)
# (0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)

由于 itertools.product 返回一个生成器，因此整个列表不会按要求存储在内存中。