You could do this recursively, but it's generally best to avoid recursion in Python unless you really need it, eg, when processing recursive data structures (like trees). Recursion in standard Python (aka CPython) is not very efficient because it cannot do tail call elimination. Also, it applies a recursion limit (which is by default 1000 levels, but that can be modified by the user).
The sequences that you want to generate are known as weak compositions, and the Wikipedia article gives a simple algorithm which is easy to implement with the help of the standard itertools.combinations
function.
#!/usr/bin/env python3
''' Generate the compositions of num of a given width
Algorithm from
https://en.wikipedia.org/wiki/Composition_%28combinatorics%29#Number_of_compositions
Written by PM 2Ring 2016.11.11
'''
from itertools import combinations
def compositions(num, width):
m = num + width - 1
last = (m,)
first = (-1,)
for t in combinations(range(m), width - 1):
yield [v - u - 1 for u, v in zip(first + t, t + last)]
# test
for t in compositions(5, 4):
print(t)
print('- ' * 20)
for t in compositions(3, 3):
print(t)
output
[0, 0, 0, 5]
[0, 0, 1, 4]
[0, 0, 2, 3]
[0, 0, 3, 2]
[0, 0, 4, 1]
[0, 0, 5, 0]
[0, 1, 0, 4]
[0, 1, 1, 3]
[0, 1, 2, 2]
[0, 1, 3, 1]
[0, 1, 4, 0]
[0, 2, 0, 3]
[0, 2, 1, 2]
[0, 2, 2, 1]
[0, 2, 3, 0]
[0, 3, 0, 2]
[0, 3, 1, 1]
[0, 3, 2, 0]
[0, 4, 0, 1]
[0, 4, 1, 0]
[0, 5, 0, 0]
[1, 0, 0, 4]
[1, 0, 1, 3]
[1, 0, 2, 2]
[1, 0, 3, 1]
[1, 0, 4, 0]
[1, 1, 0, 3]
[1, 1, 1, 2]
[1, 1, 2, 1]
[1, 1, 3, 0]
[1, 2, 0, 2]
[1, 2, 1, 1]
[1, 2, 2, 0]
[1, 3, 0, 1]
[1, 3, 1, 0]
[1, 4, 0, 0]
[2, 0, 0, 3]
[2, 0, 1, 2]
[2, 0, 2, 1]
[2, 0, 3, 0]
[2, 1, 0, 2]
[2, 1, 1, 1]
[2, 1, 2, 0]
[2, 2, 0, 1]
[2, 2, 1, 0]
[2, 3, 0, 0]
[3, 0, 0, 2]
[3, 0, 1, 1]
[3, 0, 2, 0]
[3, 1, 0, 1]
[3, 1, 1, 0]
[3, 2, 0, 0]
[4, 0, 0, 1]
[4, 0, 1, 0]
[4, 1, 0, 0]
[5, 0, 0, 0]
- - - - - - - - - - - - - - - - - - - -
[0, 0, 3]
[0, 1, 2]
[0, 2, 1]
[0, 3, 0]
[1, 0, 2]
[1, 1, 1]
[1, 2, 0]
[2, 0, 1]
[2, 1, 0]
[3, 0, 0]
FWIW, the above code can generate the 170544 sequences of compositions(15, 8)
in around 1.6 seconds on my old 2GHz 32bit machine, running on Python 3.6 or Python 2.6. (The timing information was obtained by using the Bash time
command).
FWIW, here's a recursive version taken from this answer by user3736966. I've modified it to use the same argument names as my code, to use lists instead of tuples, and to be compatible with Python 3.
def compositions(num, width, parent=[]):
if width > 1:
for i in range(num, -1, -1):
yield from compositions(i, width - 1, parent + [num - i])
else:
yield parent + [num]
Somewhat surprisingly, this one is a little faster than the original version, clocking in at around 1.5 seconds for compositions(15, 8)
.
If your version of Python doesn't understand yield from
, you can do this:
def compositions(num, width, parent=[]):
if width > 1:
for i in range(num, -1, -1):
for t in compositions(i, width - 1, parent + [num - i]):
yield t
else:
yield parent + [num]
To generate the compositions in descending order, simply reverse the range
call, i.e. for i in range(num + 1):
.
Finally, here's an unreadable one-line version. :)
def c(n, w, p=[]):
yield from(t for i in range(n,-1,-1)for t in c(i,w-1,p+[n-i]))if w-1 else[p+[n]]
Being an inveterate tinkerer, I couldn't stop myself from making yet another version. :) This is simply the original version combined with the code for combinations
listed in the itertools docs. Of course, the real itertools.combinations
is written in C so it runs faster than the roughly equivalent Python code shown in the docs.
def compositions(num, width):
r = width - 1
indices = list(range(r))
revrange = range(r-1, -1, -1)
first = [-1]
last = [num + r]
yield [0] * r + [num]
while True:
for i in revrange:
if indices[i] != i + num:
break
else:
return
indices[i] += 1
for j in range(i+1, r):
indices[j] = indices[j-1] + 1
yield [v - u - 1 for u, v in zip(first + indices, indices + last)]
This version is about 50% slower than the original at doing compositions(15, 8)
: it takes around 2.3 seconds on my machine.