CODEWITHSUNDEEP
pythonalgorithmset
Jordan
Jordan

Reputation: 133

Combinations from multiple sets where each result has only one element from each set

So I have some sets:

sets = [
   {a1,a2}, #a set
   {b1,b2}, #b set
   {c1,c2}, #c set
   ...
]

And I need to get all possible combinations of all sizes... but only one element from each set.

[
   {a1}, {b1}, ...,                        # all single
   {a1, b1}, {a1, b2} ...,                 # all pairs
   ...                                     # all combinations less than n
   {a1,b1,c1 ... }, {a2,b1,c1 ... }, ...,  # size n
]

At the moment I am using iter_tools' powerset on the union of all the sets then filtering, but this is incredibly slow when the sets get large.

large_set = union_all_sets(sets)
r = more_itertools.powerset(large_set)
r = filter_1(r)

I was wondering if there is a simple function for this that I don't know the name of? Or if anybody had any suggestions how I can avoid calling powerset on the large set?

Upvotes: 1

Views: 498

Answers (1)

Thierry Lathuille
Thierry Lathuille

Reputation: 24280

You can build combinations of 1, then 2... n of your sets, and for each of these combinations, generate the cartesian product:

from itertools import product, combinations

sets = [
   {"a1", "a2"}, #a set
   {"b1", "b2"}, #b set
   {"c1", "c2", "c3"} #c set
]

out = []

for size in range(1, len(sets)+1):
    for combination in combinations(sets, r=size):
        out.extend(set(p) for p in product(*combination))

print(out)

Output:

[{'a2'}, {'a1'}, {'b1'}, {'b2'}, {'c2'}, {'c1'}, {'c3'}, 
 {'b1', 'a2'}, {'b2', 'a2'}, {'b1', 'a1'}, {'b2', 'a1'}, {'c2', 'a2'},
 {'c1', 'a2'}, {'c3', 'a2'}, {'c2', 'a1'}, {'c1', 'a1'}, {'c3', 'a1'}, 
 {'b1', 'c2'}, {'b1', 'c1'}, {'b1', 'c3'}, {'c2', 'b2'}, {'c1', 'b2'}, 
 {'b2', 'c3'}, 
 {'b1', 'c2', 'a2'}, {'b1', 'c1', 'a2'}, {'b1', 'c3', 'a2'}, {'c2', 'b2', 'a2'}, 
 {'c1', 'b2', 'a2'}, {'b2', 'c3', 'a2'}, {'b1', 'c2', 'a1'}, {'b1', 'c1', 'a1'}, 
 {'b1', 'c3', 'a1'}, {'c2', 'b2', 'a1'}, {'c1', 'b2', 'a1'}, {'b2', 'c3', 'a1'}]

Upvotes: 2

Related Questions