1

I am trying to implement a generator function generate_tuples which yields every possible combination of numbers with a given sum.

Given a size, elements_sum so that

generate_tuples(size=3, elements_sum=4)

it would yield these tuples:

(0, 0, 4)
(0, 1, 3)
(0, 2, 2)
(0, 3, 1)
(0, 4, 0)
(1, 0, 3)
(1, 1, 2)
(1, 2, 1)
(1, 3, 0)
(2, 0, 2)
(2, 1, 1)
(2, 2, 0)
(3, 0, 1)
(3, 1, 0)
(4, 0, 0)

For now I ended up with multiple generators for each size. For instance:

def generate_tuples3(elements_sum: int):
    for x in range(elements_sum + 1):
        for y in range(elements_sum + 1):
            for z in range(elements_sum + 1):
                if sum([x, y, z]) != elements_sum:
                    continue
                yield x, y, z

This works, but it's very inefficient and includes a lot of code duplication.

Uberbaza
  • 31
  • 3
  • 1
    In combinatorics, these are called the weak compositions of n (`elements_sum`) into k (`size`) parts. I can't make time now to say more, but searching for the right term will uncover a lot for you. It's possible to generate them directly, one at a time, with an efficient loop-free algorithm (although you need a mutable representation to get away with that). – Tim Peters May 01 '22 at 16:43

1 Answers1

2

Thanks to comment posted by Tim Peters, I was able to find this algorithm which is exactly what I need!

Uberbaza
  • 31
  • 3