Generating arrays of non-negative integers where the sum of the elements is fixed

user_11688815 0
Viewed 726

Given an integer n and k, I want to create an array of all arrays of size k containing non-negative integers which sum to n. For example, if k=3 and n=10 I would want

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

Note that the order matters, which might make this easier. I know that there should be n+k-1 choose k-1 arrays in total.

My original idea was to just have k nested loops which would go through 0 to n on each element and then have an if statement at the end to check that the sum is n. This seems clumsy and very inefficient though, and I was wondering if there was a better way of doing it, ideally avoiding nested loops because I would like to be able to adjust k easily. Perhaps there is a relevant library? I am using Python.

This is what I have for k=4 and n=16 (yuck):

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
combinatorics math python
2 Answers

Here is one way, using the elegant 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

For example, sums(10,3) evaluates to:

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

Your problem can be solved by recursion. The idea is to figure out what the possibilities are for the first number in the sequence. Then for each of those possibilities, fix the first number to that value and find all the possibilities for all remaining places in the sequence. Note that I use the parameter r rather than k. In the spirit of the itertools module, this is a generator, and each yielded partition is a tuple, rather than the lists that you show. These are yielded in sorted order.

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

We can see the results from the code

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

and the printout is

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