Генерация массивов неотрицательных целых чисел, в которых сумма элементов фиксирована

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

Вот один из способов, используя элегантный трюк звезд и баров :

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