python - 生成元素之和固定的非负整数数组

Question

给定一个整数 n 和 k，我想创建一个包含所有大小为 k 的数组的数组，其中包含总和为 n 的非负整数。例如，如果 k=3 和 n=10 我想要

[2,2,6]
[2,6,2]
[3,3,4]
[10,0,0]
etc....

请注意顺序很重要，这可能会使这更容易。我知道总共应该有n+k-1个选择k-1个数组。

我最初的想法是只有 k 个嵌套循环，它会在每个元素上从 0 到 n，然后在末尾有一个 if 语句来检查总和是否为 n。这看起来很笨拙而且效率很低，我想知道是否有更好的方法，最好避免嵌套循环，因为我希望能够轻松调整 k。也许有相关的图书馆？我正在使用 Python。

这就是我对 k=4 和 n=16 的看法(糟糕):

a=0
list = []
for i in range(17):
    for j in range(17-i):
        for k in range(17-i-j):
            for l in range(17-i-j-k):
                if i+j+k+l==16:
                    list.append([i,j,k,l])
                    a += 1   

Answer 1

这是一种方法，使用优雅的 stars and bars trick :

#uses stars and bars to enumerate k-tuples of nonnegative numbers which sum to n:

import itertools

def sums(n,k):
    solutions = []
    for combo in itertools.combinations(range(n+k-1),k-1):
        s = [combo[0]]
        for i in range(1,k-1):
            s.append(combo[i]-combo[i-1]-1)
        s.append(n+k-2 - combo[k-2])
        solutions.append(s)
    return solutions

例如，sums(10,3) 的计算结果为:

[[0, 0, 10], [0, 1, 9], [0, 2, 8], [0, 3, 7], [0, 4, 6], [0, 5, 5], [0, 6, 4], [0, 7, 3], [0, 8, 2], [0, 9, 1], [0, 10, 0], [1, 0, 9], [1, 1, 8], [1, 2, 7], [1, 3, 6], [1, 4, 5], [1, 5, 4], [1, 6, 3], [1, 7, 2], [1, 8, 1], [1, 9, 0], [2, 0, 8], [2, 1, 7], [2, 2, 6], [2, 3, 5], [2, 4, 4], [2, 5, 3], [2, 6, 2], [2, 7, 1], [2, 8, 0], [3, 0, 7], [3, 1, 6], [3, 2, 5], [3, 3, 4], [3, 4, 3], [3, 5, 2], [3, 6, 1], [3, 7, 0], [4, 0, 6], [4, 1, 5], [4, 2, 4], [4, 3, 3], [4, 4, 2], [4, 5, 1], [4, 6, 0], [5, 0, 5], [5, 1, 4], [5, 2, 3], [5, 3, 2], [5, 4, 1], [5, 5, 0], [6, 0, 4], [6, 1, 3], [6, 2, 2], [6, 3, 1], [6, 4, 0], [7, 0, 3], [7, 1, 2], [7, 2, 1], [7, 3, 0], [8, 0, 2], [8, 1, 1], [8, 2, 0], [9, 0, 1], [9, 1, 0], [10, 0, 0]]