Dynamic Nesting of For loops in python

Question

I have a small piece of code that prints all possible ways of filling 'b' balls in 'c' cups considering we can also choose a subset of the balls (like no balls at all):

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)

The result is:

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

Now, the problem is that this is not a generic code and depending on the number of cups (c), the number of nested For loops needs to be changed. I tried using recursion but was not able to get the desired result.

Can anyone help me with making this code generic such that I just need to enter the value of 'b' and 'c' and get the result without editing the code?

Answer 1

Here is a solution using recursive generators

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]