Генерация всех возможных списков длины N, которые суммируются с S в Python

Question

Я пытаюсь сгенерировать все возможные списки длины N, которые суммируются с S. Я написал некоторый код для этого, но для чего угодно большого (в частности, я хочу N = 5, S = 100), я возникают ошибки переполнения памяти.

Я ищу или лучшее решение проблемы, или способ улучшить мой код, чтобы я мог запустить его на N = 5, S = 100. Эти две программы, приведенные ниже, работают в тандеме, чтобы создать все возможные комбинации чисел во вложенных списках, а затем переработать их в нужный формат. Некоторые примеры выходных данных воспроизводятся ниже.

Я знаю, что мой код не самый лучший. Я по профессии инженер (я знаю, я знаю), поэтому кодирование не совсем моя сильная сторона. Я ценю любую помощь, которую вы можете оказать.

РЕДАКТИРОВАТЬ: Я просто хотел уточнить кое-что. Во-первых, в списках нормально иметь ноль, списки могут содержать несколько одинаковых номеров, и порядок чисел в списках имеет значение.

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

Использование генератора экономит память (используйте xrange вместо range при использовании Python 2).Это то, что я придумал.Это очень похоже на ваш nToSum без необходимости 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)

Вывод

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)