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

Question

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

Lists of integers which sum to a given value $n$ are a called "integer partitions" of $n$. Searching for this term gives more solutions, *e.g.* [here](https://stackoverflow.com/questions/10035752/elegant-python-code-for-integer-partitioning). — jochen, Jun 23 '19 at 17:53

John Coleman · Accepted Answer · 2019-06-23T23:22:38.620

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

@Colonel In the final step of the decoding process I didn't use the right offset. I used `n+1` but should have used `n+k-2`. When `k = 3` (my test case), these two expressions agree. I should have checked other cases. Thanks for pointing it out. It should work now. — John Coleman, Jun 23 '19 at 23:25

Rory Daulton · Answer 2 · 2019-06-23T17:43:27.627

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)

It is interesting to note how our solutions, while being very different in some ways, nevertheless enumerate the solutions in the same order. — John Coleman, Jun 23 '19 at 17:45
@JohnColeman; That order was deliberate. It comes most naturally from the algorithms, and lexicographical order makes it easy to find a desired partition. This also follows the spirit of the `itertools` module, which you used and I imitated. — Rory Daulton, Jun 23 '19 at 17:49

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

2 Answers2