Generate all set partitions with restrictions

Given N arrays with sizes > 0. Generate all set partitions with restrictions.

Restrictions are defined as prohibiting combinations of elements from the same array.

Example:

// Given arrays
A = [A1, A2, A3]
B = [B1, B2]
C = [C1, C2]

Allowed partition:

{[A1,B1],[B2],[A2,C1],[A3,C2]}

Prohibited partition:

{[A1, B1, B2],[A2,C1],[A3,C2]} // B1, B2 from same array can't be grouped

Combining the combinatorial tools from python's itertools module:

from itertools import chain, combinations, permutations

def powerset(iterable):
    "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))

def partitions2(A, B):
    A, B = map(set, (A, B))
    for A_paired in powerset(A):
        A_alone = A.difference(A_paired)
        for B_paired in permutations(B, len(A_paired)):
            B_alone = B.difference(B_paired)
            yield tuple(chain(zip(A_paired, B_paired), zip(A_alone), zip(B_alone)))

print( list(partitions2('ab', (0,1))) )
# [(('b',), ('a',), (0,), (1,)),
#  (('b', 0), ('a',), (1,)),
#  (('b', 1), ('a',), (0,)),
#  (('a', 0), ('b',), (1,)),
#  (('a', 1), ('b',), (0,)),
#  (('b', 0), ('a', 1)),
#  (('b', 1), ('a', 0))]

def partitions(*iterables):
    if len(iterables) == 0:
        yield ((),)
    elif len(iterables) == 1:
        yield tuple(zip(iterables[0]))
    elif len(iterables) == 2:
        yield from partitions2(*iterables)
    else:
        A, *BCD = iterables
        A = set(A)
        for P in map(set,partitions(*BCD)):
            for A_paired in powerset(A):
                A_alone = A.difference(A_paired)
                for P_paired in permutations(P, len(A_paired)):
                    P_alone = P.difference(P_paired)
                    yield tuple(chain(((a, *p) for a,p in zip(A_paired, P_paired)), zip(A_alone), P_alone))

print( list(partitions('x', 'a', (0,1))) )
# [(('x',), (0,), (1,), ('a',)),
#  (('x', 0), (1,), ('a',)),
#  (('x', 1), (0,), ('a',)),
#  (('x', 'a'), (0,), (1,)),
#  (('x',), (1,), ('a', 0)),
#  (('x', 1), ('a', 0)),
#  (('x', 'a', 0), (1,)),
#  (('x',), (0,), ('a', 1)),
#  (('x', 0), ('a', 1)),
#  (('x', 'a', 1), (0,))]

The number of partitions explodes quite fast. Here is the specific example you asked for:

print( list(partitions(('A1', 'A2', 'A3'), ('B1', 'B2'), ('C1', 'C2'))) )
# [(('A2',), ('A3',), ('A1',), ('B2',), ('C1',), ('B1',), ('C2',)),
#  (('A2', 'B2'), ('A3',), ('A1',), ('C1',), ('B1',), ('C2',)),
#  (('A2', 'C1'), ('A3',), ('A1',), ('B2',), ('B1',), ('C2',)),
#  (('A2', 'B1'), ('A3',), ('A1',), ('B2',), ('C1',), ('C2',)),
#  (('A2', 'C2'), ('A3',), ('A1',), ('B2',), ('C1',), ('B1',)),
#  (('A3', 'B2'), ('A2',), ('A1',), ('C1',), ('B1',), ('C2',)),
#  (('A3', 'C1'), ('A2',), ('A1',), ('B2',), ('B1',), ('C2',)),
# ...
#  (('A2', 'C1'), ('A3', 'C2'), ('A1', 'B1'), ('B2',)),  # <----- THE EXAMPLE YOU GAVE
#  (('A2', 'B1'), ('A3', 'B2'), ('A1', 'C1'), ('C2',)),
#  (('A2', 'B1'), ('A3', 'B2'), ('A1', 'C2'), ('C1',)),
#  (('A2', 'B1'), ('A3', 'C1'), ('A1', 'B2'), ('C2',)),
#  (('A2', 'B1'), ('A3', 'C1'), ('A1', 'C2'), ('B2',)),
# ... 
#  (('A2', 'B1', 'C1'), ('A1', 'B2', 'C2'), ('A3',)),
#  (('A3', 'B2', 'C2'), ('A1', 'B1', 'C1'), ('A2',)),
#  (('A3', 'B1', 'C1'), ('A1', 'B2', 'C2'), ('A2',)),
#  (('A2',), ('A3',), ('A1',), ('B1', 'C2'), ('B2', 'C1')),
#  (('A2', 'B1', 'C2'), ('A3',), ('A1',), ('B2', 'C1')),
#  (('A2', 'B2', 'C1'), ('A3',), ('A1',), ('B1', 'C2')),
# ...
#  (('A3', 'B1', 'C2'), ('A1', 'B2', 'C1'), ('A2',)),
#  (('A3', 'B2', 'C1'), ('A1', 'B1', 'C2'), ('A2',))]

We can get the powerset of the different arrays, to get all possible combinations of arrays, and then take the cartesian product of those arrays, to form the actual resulting sets.

Taking the implementation of powerset from the python itertools docs:

from itertools import chain, combinations, product

def powerset(iterable):
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))

def solve(*arrays):
    for subset in powerset(arrays):
        yield from product(*subset)

Example:

>>> list(solve(["a1", "a2", "a3"], ["b1", "b2"], ["c1", "c2"]))
[(), ('a1',), ('a2',), ('a3',), ('b1',), ('b2',), ('c1',), ('c2',), ('a1', 'b1'), ('a1', 'b2'), ('a2', 'b1'), ('a2', 'b2'), ('a3', 'b1'), ('a3', 'b2'), ('a1', 'c1'), ('a1', 'c2'), ('a2', 'c1'), ('a2', 'c2'), ('a3', 'c1'), ('a3', 'c2'), ('b1', 'c1'), ('b1', 'c2'), ('b2', 'c1'), ('b2', 'c2'), ('a1', 'b1', 'c1'), ('a1', 'b1', 'c2'), ('a1', 'b2', 'c1'), ('a1', 'b2', 'c2'), ('a2', 'b1', 'c1'), ('a2', 'b1', 'c2'), ('a2', 'b2', 'c1'), ('a2', 'b2', 'c2'), ('a3', 'b1', 'c1'), ('a3', 'b1', 'c2'), ('a3', 'b2', 'c1'), ('a3', 'b2', 'c2')]

You can of course skip the empty set () if you like.

Counting Update

If you just want the count of "partitions", then the solution is much easier. You could still of course use the above generator, convert it to a list, and get its len, but this will be much faster:

from math import prod

def count(*arrays):
    return prod(len(array)+1 for array in arrays)

Note that math.prod requires python 3.8+

For the example:

>>> count(["a1", "a2", "a3"], ["b1", "b2"], ["c1", "c2"])
36