generating conditional product of lists in python (combinatorics)

Question

I want to be able to generate a conditional product. So similar to this answer: All combinations of a list of lists

I wanted to use itertools.product(*listOfLists). However, my problem is that the inclusion of one element from one list means other lists must be consulted for the product.

Example:

colors = ['red', 'blue', 'green']
fruits = ['apple', 'orange', 'banana']
locations = ['indoors', 'outdoors']

indoor_choices = ['bathroom', 'bedroom', 'kitchen']
green_choices = ['forest', 'light', 'dark']

Here, we want to consider always every possible choice of color, fuit, and location. However, in the case of 'indoor', we also want to consider indoor_choices, and in the case of 'green' being in a possible choice, we also want to choose a more specific color of green. It's kind of a tree of possibilities where some branches keep branching and others do not.

So in this silly example above you could do a for loop like so:

for c in colors:
    for f in fruits:
        for l in locations:
            # etc

but then we encounter the problem of what happens when two different categories have possible branching based on this choice.

A simple (hacky) solution would be just to manually code conditions and put for loops inside of them:

for c in colors:
    for f in fruits:
        for l in locations:

            if c == 'green' and l == 'indoor':
                for gc in green_choices:
                     for ic in indoor_choices:
                         # output

            elif c == 'green':
                for gc in green_choices:
                    # output

            elif l == 'indoor':
                for gc in green_choices:
                    # output

            else:
                # output

but imagine the horror when there are N lists where M of them have additional branching. Or even worse, there is nested additional branching ... basically this hack doesn't scale.

Any ideas? This problem has proved itself to be deceptively hard!

Answer 1

Here's how I'd do it, with a recursive generator.

def prod(terms, expansions):
    if not terms: # base case
        yield ()
        return

    t = terms[0] # take the first term

    for v in expansions[t]: # expand the term, to get values
        if v not in expansions: # can the value can be expanded?
            gen = prod(terms[1:], expansions) # if not, we do a basic recursion
        else:
            gen = prod(terms[1:] + [v], expansions) # if so, we add it to terms

        for p in gen: # now we get iterate over the results of the recursive call
            yield (v,) + p # and add our value to the start

Here's how you call it to generate the product you wanted in your example:

expansions = {
        'colors':['red', 'blue', 'green'],
        'fruits':['apple', 'orange', 'banana'],
        'locations':['indoors', 'outdoors'],
        'indoors':['bathroom', 'bedroom', 'kitchen'],
        'green':['forest', 'light', 'dark']
    }

terms = ["colors", "locations"] # fruits omitted, to reduce the number of lines

for p in prod(terms, expansions):
    print(p)

Output:

('red', 'indoors', 'bathroom')
('red', 'indoors', 'bedroom')
('red', 'indoors', 'kitchen')
('red', 'outdoors')
('blue', 'indoors', 'bathroom')
('blue', 'indoors', 'bedroom')
('blue', 'indoors', 'kitchen')
('blue', 'outdoors')
('green', 'indoors', 'forest', 'bathroom')
('green', 'indoors', 'forest', 'bedroom')
('green', 'indoors', 'forest', 'kitchen')
('green', 'indoors', 'light', 'bathroom')
('green', 'indoors', 'light', 'bedroom')
('green', 'indoors', 'light', 'kitchen')
('green', 'indoors', 'dark', 'bathroom')
('green', 'indoors', 'dark', 'bedroom')
('green', 'indoors', 'dark', 'kitchen')
('green', 'outdoors', 'forest')
('green', 'outdoors', 'light')
('green', 'outdoors', 'dark')

Answer 2

If your real problem is really just like your example, then you can analyze the combinations into just four products:

is_green = ['green']
not_green = ['red', 'blue']
is_indoors = ['indoors']
not_indoors = ['outdoors']

p1 = itertools.product([not_green, fruits, not_indoors])
...
p2 = itertools.product([is_green, fruits, not_indoors, green_choices])
...
p3 = itertools.product([not_green, fruits, is_indoors, indoor_choices])
...
p4 = itertools.product([is_green, fruits, is_indoors, green_choices, indoor_choices])

That's all!

Now if we want to generalize so we don't have to make four "special" cases, we can capture the relation between certain values and the additional choices that they open up, as @DavidRobinson suggested.

import itertools

colors = ['red', 'blue', 'green']
fruits = ['apple', 'orange', 'banana']
locations = ['indoors', 'outdoors']

indoor_choices = ('bathroom', 'bedroom', 'kitchen')
green_choices = ('forest', 'light', 'dark')

choices = [colors, fruits, locations]
more_choices = { 'indoors': indoor_choices, 'green': green_choices }
for p in itertools.product(*choices):
    m = [more_choices[k] for k in p if k in more_choices]
    for r in itertools.product([p],*m):
        print list(r[0]) + list(r[1:])

Please note that there will inevitably be difficulties when choices and more_choices are large.