生成元素总和固定的非负整数数组

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

Answer 2

你的问题可以通过递归来解决。这个想法是找出序列中第一个数字的可能性。然后对于这些可能性中的每一个，将第一个数字固定为该值并找到序列中所有剩余位置的所有可能性。请注意，我使用参数r 而不是k。本着itertools 模块的精神，这是一个生成器，每个产生的分区都是一个元组，而不是您显示的列表。这些是按排序顺序生成的。

def partitions_nonnegative_fixed_length_ordered(n, r):
    """Generate the partitions of the nonnegative integer `n` as the
    sum of `r` nonnegative integers, where the order of the integers
    matters. The partitions are tuples and are generated in
    lexicographic order. The number of partitions generated is
    binomialcoefficient(n+r-1, r-1).

    NOTE:   The empty generator is returned for n=r=0, though arguably
            the generator yielding just the empty tuple would satisfy
            the conditions.
    """
    def partitions_prefixed(prefix, n, r):
        if r == 1:
            yield prefix + (n,)
        else:
            for i in range(n + 1):
                yield from partitions_prefixed(prefix + (i,), n - i, r - 1)

    if n >= 0 and r >= 1 and n == int(n) and r == int(r):
        yield from partitions_prefixed(tuple(), int(n), int(r))

我们可以从代码中看到结果

for partition in partitions_nonnegative_fixed_length_ordered(10, 3):
    print(partition)

打印出来的是

(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)