1

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:

  1. Generate all bit-arrays of size 2n-1 and n bits on (using itertools.combinations)
  2. Regard zeros as partitions and ones as balls - sum the number of ones between two zeros (with reduce)
  3. Sorting all options, and removing duplicates with set()

The code works, but it is very slow.

Any ideas how can it be improved ?

Uri Goren
  • 13,386
  • 6
  • 58
  • 110

0 Answers0