在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 c in colors:
    for f in fruits:
        for l in locations:
            # etc

但这样我们就会遇到一个问题：当两个不同的类别根据这个选择有可能的分支时，会发生什么。

一个简单（但有点笨拙）的解决方案就是手动编写条件，并把循环放在里面：

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

如果你的实际问题和你给的例子差不多，那你可以把组合分析成四种产品：

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

就这样！

现在如果我们想要更通用一点，这样就不需要处理四个“特殊”情况了，我们可以抓住某些值之间的关系，以及它们所带来的额外选择，正如@DavidRobinson所建议的那样。

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:])

请注意，当选择和更多选择的数量很大时，难免会遇到一些困难。

Answer 2

这里有一个使用 yield 的递归实现。我觉得它没有 @Blckknght 的解决方案那么简洁，但可能对你有帮助。

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

green_subtypes = ["forest", "light", "dark"]
indoor_locations = ["bathroom","bedroom","kitchen"]

def gen(state):
  if len(state)==0:
    for c in colors:
       s = [c]
       for y in gen(s):
         yield y
  elif len(state)==1:
    for x in fruits:
      s = state + [x]
      for y in gen(s):
        yield y
  elif len(state)==2:
    for x in locations:
      s = state + [x]
      for y in gen(s):
        yield y
  else:
    # If we're green and we haven't looped through the green options already 
    # (the check is a bit dodgy and could do with being moved into a flag inside state)
    if state[0]=='green' and len(set(state).intersection(set(green_subtypes)))==0:
      for x in green_subtypes:
        s = state + [x]
        for y in gen(s):
          yield y
    # If we're indoors and we haven't looped through the indoor options already 
    # (the check is a bit dodgy and could do with being moved into a flag inside state)
    elif state[2]=='indoors' and len(set(state).intersection(set(indoor_locations)))==0:
      for x in indoor_locations:
        s = state + [x]
        for y in gen(s):
          yield y
    else:
      yield state

for x in gen([]):
  print(x)

Answer 3

这是我会做的，使用一个递归生成器。

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