Given 3 balls, there are 3 ways to partition them into bins:
(3)
(2,1)
(1,1,1)
Note that (1,2)
is already covered in (2,1)
, since we don't care about the order of the bins.
I'd like to write a function that gets the number of balls n
and outputs all possible partitions, in an efficient way.
This is what I came up with:
from itertools import combinations
from functools import reduce
def all_partitions(n):
n_bits_on=set([tuple(reduce(lambda x,y:x[:y]+[1]+x[y+1:],c,[0]*(2*n-1))) for c in combinations(range(2*n-1),n)])
return set([tuple(sorted(filter(lambda i:i>0,reduce(lambda x,y: x+[y] if y==0 else x[:-1]+[x[-1]+1,],p,[0])))) for p in n_bits_on])
What the code does is:
- Generate all bit-arrays of size
2n-1
andn
bits on (usingitertools.combinations
) - Regard zeros as partitions and ones as balls - sum the number of ones between two zeros (with
reduce
) - Sorting all options, and removing duplicates with
set()
The code works, but it is very slow.
Any ideas how can it be improved ?