在Python中生成所有长度为N且和为S的可能列表

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

我在找类似的问题，但想要在每个选项中增加一些限制条件。以下是我的方法：

比如说，如果我们想要从10到50这5个数字中找出一些排列，使它们的和等于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)

Answer 2

我发现这里的答案非常有帮助。我添加了一个版本，使用了一个范围内的增量，因为这符合我的需求。注意：需要使用NumPy来支持非整数的增量。同时假设0是默认的最小值。

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

使用生成器可以节省内存（如果你在用Python 2，记得用xrange代替range）。这是我想到的解决方案。它和你的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)