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')