General bars and stars

Question

I've got a the following "bars and stars" algorithm, implemented in Python, which prints out all decomposition of a sum into 3 bins, for sums going from 0 to 5. I'd like to generalise my code so it works with N bins (where N less than the max sum i.e 5 here). The pattern is if you have 3 bins you need 2 nested loops, if you have N bins you need N-1 nested loops.

Can someone think of a generic way of writing this, possibly not using loops?

# bars and stars algorithm
N=5
for n in range(0,N):
    x=[1]*n
    for i in range(0,(len(x)+1)):
        for j in range(i,(len(x)+1)):
            print sum(x[0:i]), sum(x[i:j]), sum(x[j:len(x)])

Answer 1

If this isn't simply a learning exercise, then it's not necessary for you to roll your own algorithm to generate the partitions: Python's standard library already has most of what you need, in the form of the itertools.combinations function.

From Theorem 2 on the Wikipedia page you linked to, there are n+k-1 choose k-1 ways of partitioning n items into k bins, and the proof of that theorem gives an explicit correspondence between the combinations and the partitions. So all we need is (1) a way to generate those combinations, and (2) code to translate each combination to the corresponding partition. The itertools.combinations function already provides the first ingredient. For the second, each combination gives the positions of the dividers; the differences between successive divider positions (minus one) give the partition sizes. Here's the code:

import itertools

def partitions(n, k):
    for c in itertools.combinations(range(n+k-1), k-1):
        yield [b-a-1 for a, b in zip((-1,)+c, c+(n+k-1,))]

# Example usage
for p in partitions(5, 3):
    print(p)

And here's the output from running the above code.

[0, 0, 5]
[0, 1, 4]
[0, 2, 3]
[0, 3, 2]
[0, 4, 1]
[0, 5, 0]
[1, 0, 4]
[1, 1, 3]
[1, 2, 2]
[1, 3, 1]
[1, 4, 0]
[2, 0, 3]
[2, 1, 2]
[2, 2, 1]
[2, 3, 0]
[3, 0, 2]
[3, 1, 1]
[3, 2, 0]
[4, 0, 1]
[4, 1, 0]
[5, 0, 0]