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]

Another recursive variant, using a generator function, i.e. instead of right away printing the results, it yields them one after another, to be printed by the caller.

The way to convert your loops into a recursive algorithm is as follows:

identify the "base case": when there are no more bars, just print the stars
for any number of stars in the first segment, recursively determine the possible partitions of the rest, and combine them

You can also turn this into an algorithm to partition arbitrary sequences into chunks:

def partition(seq, n, min_size=0):
    if n == 0:
        yield [seq]
    else:
        for i in range(min_size, len(seq) - min_size * n + 1):
            for res in partition(seq[i:], n-1, min_size):
                yield [seq[:i]] + res

Example usage:

for res in partition("*****", 2):
    print "|".join(res)

Take it one step at a time.

First, remove the sum() calls. We don't need them:

N=5
for n in range(0,N):
    x=[1]*n
    for i in range(0,(n+1)):  # len(x) == n
        for j in range(i,(n+1)):
            print i, j - i, n - j

Notice that x is an unused variable:

N=5
for n in range(0,N):
    for i in range(0,(n+1)):
        for j in range(i,(n+1)):
            print i, j - i, n - j

Time to generalize. The above algorithm is correct for N stars and three bars, so we just need to generalize the bars.

Do this recursively. For the base case, we have either zero bars or zero stars, which are both trivial. For the recursive case, run through all the possible positions of the leftmost bar and recurse in each case:

from __future__ import print_function

def bars_and_stars(bars=3, stars=5, _prefix=''):
    if stars == 0:
        print(_prefix + ', '.join('0'*(bars+1)))
        return
    if bars == 0:
        print(_prefix + str(stars))
        return
    for i in range(stars+1):
        bars_and_stars(bars-1, stars-i, '{}{}, '.format(_prefix, i))

For bonus points, we could change range() to xrange(), but that will just give you trouble when you port to Python 3.

General bars and stars

Related

Recent Posts