在python中生成列表的条件乘积(组合)

Question

我希望能够生成有条件的产品.与此答案类似: 列表列表的所有组合

我想用itertools.product(*listOfLists).但是,我的问题是从一个列表中包含一个元素意味着必须为产品查阅其他列表.

例:

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

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

在这里,我们要始终考虑每种可能的颜色,特性和位置选择.然而,在"室内"的情况下,我们也想考虑室内选择,并且在"绿色"可能的选择的情况下,我们还想选择更具体的绿色.它是一种可能性的树,其中一些分支保持分支而另一些则不分支.

所以在上面的这个愚蠢的例子中你可以这样做一个for循环:

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

但是,当我们遇到两个不同类别根据这个选择进行分支时会发生什么的问题.

一个简单的(hacky)解决方案就是手动编写条件并在其中放置循环:

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

但想象一下,如果有N个列表,其中M个有额外的分支,那就是恐怖.或者更糟糕的是,有嵌套的额外分支...基本上这个黑客不会扩展.

有任何想法吗？事实证明这个问题看起来很难!

Answer 1

以下是我使用递归生成器的方法.

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

以下是您在示例中调用它来生成所需产品的方法:

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)

输出:

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