How to generate a set of all tuples of given length and sum of elements?

Question

I would like to have a function that generates a set (or a list) of all possible tuples with a given length and sum of their elements. The elements of tuples should be not negative integer.

For example for the following input

get_tuple(length=3, total=2)

I would like to get the following output:

[(1, 0, 1), (2, 0, 0), (1, 1, 0), (0, 0, 2), (0, 1, 1), (0, 2, 0)]

Is the a standard library in Python that can do that? If not, how to write a function that can do it?

Answer 1

You can create a recursive function, avoiding creating all candidate combinations and checking their sums:

def get_tuples(length, total):
    if length == 1:
        yield (total,)
        return

    for i in xrange(total + 1):
        for t in get_tuples(length - 1, total - i):
            yield (i,) + t

If we test:

>>> list(get_tuples(3, 2))
[(0, 0, 2), (0, 1, 1), (0, 2, 0), (1, 0, 1), (1, 1, 0), (2, 0, 0)]