Python itertools - Create only a subset of all possible products

Question

I am working on a script that uses itertools to generate desired combinations of given parameters. I am having trouble generating the appropriate 'sets' however when some variables are linked, i.e. excluding certain combinations. Consider the following:

import itertools

A = ['a1','a2','a3']
B = ['b1','b2','b3']
C = ['c1','c2']

If I want to generate all possible combinations of these elements, I can simply use itertools.product()

all_combinations = list(itertools.product(A,B,C))

Which gives the expected

[('a1', 'b1', 'c1'), ('a1', 'b1', 'c2'), ('a1', 'b2', 'c1'), ... 
('a3', 'b2', 'c2'), ('a3', 'b3', 'c1'), ('a3', 'b3', 'c2')]

For 18 combinations (3*3*2)

However, how can I 'link' parameters A and B such that each returned set contains only 'an','bn' elements? That is, I have tried:

ABprecombine = zip(A,B)
limited_combinations = list(itertools.product(ABprecombine,C))

Which returns

[(('a1', 'b1'), 'c1'), (('a1', 'b1'), 'c2'), (('a2', 'b2'), 'c1'),
 (('a2', 'b2'), 'c2'), (('a3', 'b3'), 'c1'), (('a3', 'b3'), 'c2')]

This is the six (3*1*2) desired products, but obviously due to the way I created it I now have an extra tuple.

Of course I could generate all combinations and then filter out given ones, but is there a smart way to 'link' parameters as above?

Answer 1

Here, zipping A and B is the right way to go. You can flatten the tuples pretty easily if you want:

limited_combinations = [(a, b, c) for ((a, b), c) in itertools.product(zip(A, B), C)]

If you want more detailed control of what combinations get produced, things can rapidly get more complicated, up to the difficulty of needing to solve NP-hard problems like boolean satisfiability. If that happens, look into existing libraries for that kind of thing.