python 中 For 循环的动态嵌套

Question

考虑到我们还可以选择球的一个子集（就像根本没有球一样），我有一小段代码可以打印在“c”杯中填充“b”球的所有可能方式：

b = 2
c = 3
for i in range(b+1):
    for k in range(b+1):
        for j in range(b+1):
            if i+j+k<=b:
                print(i,k,j)

结果是：

0 0 0
0 0 1
0 0 2
0 1 0
0 1 1
0 2 0
1 0 0
1 0 1
1 1 0
2 0 0

现在，问题在于这不是通用代码，并且根据杯子的数量（c），嵌套的 For 循环的数量需要更改。 我尝试使用递归，但无法获得所需的结果。

谁能帮我使这个代码通用，这样我只需要输入'b'和'c'的值并在不编辑代码的情况下获得结果吗？

Answer 1

您不需要递归，只需itertools.product ：

for t in product(range(b+1), repeat=c):
    if sum(t) <= b:
        print(t)

但是，递归将通过动态调整每个范围以考虑先前的选择来减少您需要考虑的元组的数量。

# Untested
def foo(b, c, stop=0):
    for x in range(0, b+1-stop):
        for y in foo(b, c-1, stop + x):
            yield (x, *y)

for t in foo(b, c):
    if sum(t) <= b:
        print(t)

Answer 2

这是使用递归生成器的解决方案

def fill(b, c):
    if c == 0:
        yield []
    else:
        for i in range(b+1):
            for sub in fill(b-i, c-1):
                yield [i, *sub]



for o in fill(2, 3):
    print(o)

Output

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