Generate all possible lists of length N that sum to S in Python

Question

I'm trying to generate all possible lists of Length N that sum to S. I've written some code to do so, but on anything large (in particular, I want N=5, S=100), I run into memory overflow errors.

I'm looking for either a better solution to the problem, or a way to improve my code so I can run it on N=5, S=100. These two programs below work in tandem to create all the possible number combinations in nested lists, and then re-work them to be the right format. Some sample output is reproduced below.

I know my code isn't the best. I'm an engineer by trade (I know, I know), so coding isn't exactly my forte. I appreciate any help you can provide.

EDIT: I just wanted to clarify a few thing. First, it's ok to have zero in the lists, the lists can contain multiples of the same number, and the order of the numbers in the lists matters.

def nToSum(N,S):
    ''' Creates a nested list of all possible lists of length N that sum to S'''
    if N <= 1: #base case
        return [S]
    else:
        L = []
        for x in range(S+1):   #create a sub-list for each possible entry of 0 to S 
            L += [[x,nToSum(N-1,S-x)]]  #create a sub-list for this value recursively
        return L

def compress(n=[],L): #designed to take in a list generated by nToSum
    '''takes the input from nToSum as list L, and then flattens it so that each list is a
       top level list.  Leading set n is the "prefix" list, and grows as you climb down the 
       sublists'''
    if type(L[0]) == int:  #base case:  you have exposed a pure integer
        return [n+L]       #take that integer, and prepend the leading set n
    else:
        Q = []
        for x in L:  # look at every sublist
            Q += compress(n+[x[0]],x[1])  # for each sublist, create top level lists recursively
        return Q                          # note:  append x[0] to leading set n

>>> nToSum(3,3)
[[0, [[0, [3]], [1, [2]], [2, [1]], [3, [0]]]], [1, [[0, [2]], [1, [1]], [2, [0]]]], [2, [[0, [1]], [1, [0]]]], [3, [[0, [0]]]]]

>>> compress([],nToSum(3,3))
[[0, 0, 3], [0, 1, 2], [0, 2, 1], [0, 3, 0], [1, 0, 2], [1, 1, 1], [1, 2, 0], [2, 0, 1], [2, 1, 0], [3, 0, 0]]

Answer 1

Using a generator saves on memory (usexrange instead of range if using Python 2). This is what I came up with. It is very similar to your nToSum without the need for compress.

def sums(length, total_sum):
    if length == 1:
        yield (total_sum,)
    else:
        for value in range(total_sum + 1):
            for permutation in sums(length - 1, total_sum - value):
                yield (value,) + permutation

L = list(sums(5,100))
print('total permutations:',len(L))

# First and last 10 of list
for i in L[:10] + L[-10:]:
    print(i)

Output

total permutations: 4598126
(0, 0, 0, 0, 100)
(0, 0, 0, 1, 99)
(0, 0, 0, 2, 98)
(0, 0, 0, 3, 97)
(0, 0, 0, 4, 96)
(0, 0, 0, 5, 95)
(0, 0, 0, 6, 94)
(0, 0, 0, 7, 93)
(0, 0, 0, 8, 92)
(0, 0, 0, 9, 91)
(98, 0, 2, 0, 0)
(98, 1, 0, 0, 1)
(98, 1, 0, 1, 0)
(98, 1, 1, 0, 0)
(98, 2, 0, 0, 0)
(99, 0, 0, 0, 1)
(99, 0, 0, 1, 0)
(99, 0, 1, 0, 0)
(99, 1, 0, 0, 0)
(100, 0, 0, 0, 0)

Answer 2

I found the answers here incredibly helpful. Adding a version with an increment used in the range because it fit my use case. Note: needed to use NumPy to allow for non-integer increments. Also assumed that 0 was the default minimum value.

import numpy as np

def permutations_w_increment(var_count, total_sum, increment):
if var_count == 1:
    yield (total_sum,)
else:
    for value in np.arange(0,total_sum + 1,increment):
        for permutation in permutations_w_increment(var_count - 1, total_sum - value, increment):
            yield (value,) + permutation

L = list(permutations_w_increment(5,100,1))
print('total permutations:',len(L))

# First and last 10 of list
for i in L[:10] + L[-10:]:
    print(i)

Answer 3

I was looking for something similar to this question, but with the added complexity of constraints in the options for each n of N. Here is my approach:

For example if we are looking for permutations of 5 numbers from 10 to 50 that add up to 100

def permutations_w_constraints(n_perm_elements, sum_total, min_value, max_value):
    # base case
    if n_perm_elements == 1:
        if (sum_total <= max_value) & (sum_total >= min_value):
            yield (sum_total,)
    else:
        for value in range(min_value, max_value + 1):
            for permutation in permutations_w_constraints(
                n_perm_elements - 1, sum_total - value, min_value, max_value
            ):
                yield (value,) + permutation

results = list(permutations_w_constraints(5, 100, 10, 50))

print('total permutations:',len(results))
for i in results[:10] + results[-10:]:
    print(i)

total permutations: 312676
(10, 10, 10, 20, 50)
(10, 10, 10, 21, 49)
(10, 10, 10, 22, 48)
(10, 10, 10, 23, 47)
(10, 10, 10, 24, 46)
(10, 10, 10, 25, 45)
(10, 10, 10, 26, 44)
(10, 10, 10, 27, 43)
(10, 10, 10, 28, 42)
(10, 10, 10, 29, 41)
(50, 18, 10, 10, 12)
(50, 18, 10, 11, 11)
(50, 18, 10, 12, 10)
(50, 18, 11, 10, 11)
(50, 18, 11, 11, 10)
(50, 18, 12, 10, 10)
(50, 19, 10, 10, 11)
(50, 19, 10, 11, 10)
(50, 19, 11, 10, 10)
(50, 20, 10, 10, 10)