0

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?

Burhan
  • 15
  • 4

2 Answers2

4

You don't need recursion, just itertools.product:

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

However, recursion would reduce the number of tuples you need to consider by dynamically adjusting each range to account for previous choices.

# 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)
chepner
  • 497,756
  • 71
  • 530
  • 681
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]
Hymns For Disco
  • 7,530
  • 2
  • 17
  • 33