Write a function that takes two natural numbers and as inputs and returns the set of all tuples of size that sum to

user_13612910 0
Viewed 606

In python 3 I am trying to write a function that takes two natural numbers and as inputs and returns the set of all tuples of size that sum to .

I have built the following functions:

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


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

which with, e.g., list(get_tuples(2, 8)) returns the correct results, and:

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

    comp = []
    for i in range(total + 1):
        for t in get_tuples(length - 1, total - i):
            comp.append((i,) + t)
        return comp

which however returns a wrong result (i.e. a single tuple). Can anybody please explain why, and how to fix the second function?

combinatorics python
3 Answers

Simply place the return one indent out:

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

    comp = []
    for i in range(total + 1):
        for t in get_tuples(length - 1, total - i):
            comp.append((i,) + t)
    return comp

Edit: Make sure you understand why

user_40994130

owninggreendragsdude's answer is the right one, because it shows exactly the one thing you need to change in your code to make it work.

However, if you were wondering how to get rid of the recursion and make the code faster at the same time, here's a solution:

def get_tuples_optimized(length, total):
    comp = []

    # Stack of partial tuples to process, initialized with an empty tuple
    todo = [(total, tuple())]

    # Loop as long as we have partial tuples on our stack
    while todo:
        # Take the next partial tuple
        amount_left, partial_tuple = todo.pop()

        if len(partial_tuple) == length - 1:
            # Put all that is left as the last element
            comp.append(partial_tuple + (amount_left,))
        else:
            # Add all possible extensions-by-one to our stack
            for i in range(amount_left + 1):
                extended_tuple = partial_tuple + (i,)
                todo.append((amount_left - i, extended_tuple))

    return comp
user_3202260

Here you go:

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

comp = []

for i in range(total + 1):
    for t in get_tuples(length - 1, total - i):
        comp.append((i,) + t)
return comp

Take the return statement out of the loop.

user_43506410
Related