生成非负整数数组,其中元素之和是固定的

ach*_*zig 3 python math combinatorics

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

[2,2,6]
[2,6,2]
[3,3,4]
[10,0,0]
etc....
Run Code Online (Sandbox Code Playgroud)

请注意,顺序很重要,这可能会使此过程变得容易。我知道总共应该有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   
Run Code Online (Sandbox Code Playgroud)

Joh*_*man 5

这是使用优雅的星星和小技巧的一种方法:

#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
Run Code Online (Sandbox Code Playgroud)

例如,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]]
Run Code Online (Sandbox Code Playgroud)