General bars and stars

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]

Leave a Comment