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

Question

我想生成所有长度为N且总和为S的可能列表。我编写了一些代码来实现这一点，但在处理任何大型数据时（特别是当N=5，S=100时），我遇到了内存溢出错误。

我正在寻找更好的解决方案或改进我的代码以便能够处理N=5，S=100。以下两个程序可以联合创建嵌套列表中的所有可能数字组合，然后重新设计它们以符合正确格式。下面是一些示例输出。

我知道我的代码并不是最好的。我是一名工程师（我知道，我知道），所以编码并不是我的专长。我感谢您提供的任何帮助。

编辑：我只想澄清几件事情。首先，在列表中放置0是可以的，列表中可以包含多个相同的数字，并且列表中数字的顺序很重要。

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]]

- Nick2253

---

你可能会发现这与集合的划分有关。 - Greg Hewgill

Answer 1

使用生成器可以节省内存（如果使用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)

- Mark Tolonen

---

非常感谢！这正是我所需要的。我喜欢双重循环的使用方式。还要感谢你向我介绍了Python中生成器的世界。我以前听说过它们，但这是我第一次看到它们真正的用途！（我想我需要更多地编码了！） - Nick2253

我在Windows 64/8GB桌面上为f(6,100)看到了内存错误。是否有一种使用字典或类似方法使其更快的方式？ - Stat-R

@Stat-R list(f(6,100)) 实例化了所有排列组合。f(6,100) 有 96,560,646 种排列组合，而六元组使用 96 字节的内存。这超过了 9GB 的内存。为了仅仅统计它们的数量，可以使用 sum(1 for i in f(6,100))。 - Mark Tolonen

如何将其转换为迭代解决方案？ - StefOverflow

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)

- Chris Hunter

Answer 3

我正在寻找类似于这个问题的东西，但是在每个N的选项中添加了约束条件的复杂性。这是我的方法：

例如，如果我们正在寻找从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)

- Nico