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

时间:2019-06-23 16:16:34

标签: python math combinatorics

给定一个整数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   

2 个答案:

答案 0 :(得分:5)

这是使用优雅的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]]

答案 1 :(得分:4)

您的问题可以通过递归解决。这样做的目的是弄清楚序列中第一个数字的可能性。然后针对每种可能性,将第一个数字固定为该值,并找到序列中所有剩余位置的所有可能性。请注意,我使用参数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)