Python 嵌套列表，所有组合

Question

我有以下形式的嵌套列表：

['A', 'B', ['a0', 'a1', 'a2'], 'C', 'D', ['a0', 'a1', 'a2'], 'E', 'F', ['a0', 'a1', 'a2']]

并且我想生成关于以下形式的子列表['a0', 'a1', 'a2']的所有组合：

['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a2']]
.             .                 .
.             .                 .
.             .                 .
['A', 'B', ['a2'], 'C', 'D', ['a2'], 'E', 'F', ['a2']]

等 总共 27 27 个列表。 我知道我必须使用itertools包，但我不知道如何使用。 欢迎任何想法。

Answer 1

您可以这样做（如果只有一个嵌套级别）：

from itertools import product

lst = ['A', 'B', ['a0', 'a1', 'a2'], 'C', 'D', ['a0', 'a1', 'a2'], 'E', 'F', ['a0', 'a1', 'a2']]


def nested_product(ls):
    lst_positions = [l for l in ls if isinstance(l, list)]
    for p in product(*lst_positions):
        it = iter(p)
        yield [e if not isinstance(e, list) else [next(it)] for e in lst]


for pr in nested_product(lst):
    print(pr)

输出（部分）

['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a2']]
['A', 'B', ['a1'], 'C', 'D', ['a0'], 'E', 'F', ['a0']]
...

使用yield from (Python 3.3+) 的替代解决方案如下：

def build_nested_list(ls, it):
    return [e if not isinstance(e, list) else [next(it)] for e in ls]


def nested_product(ls):
    lst_positions = (li for li in ls if isinstance(li, list))
    yield from (build_nested_list(ls, iter(p)) for p in product(*lst_positions))

Answer 2

如果只有一层嵌套：

import itertools

def combos(lst):
    # prep a copy without sublists so a little less can be copied
    copy = []
    lindices = []
    for i, v in enumerate(lst):
        if isinstance(v, list):
            # if it's a list we note the index
            lindices.append(i)
            # we append None to the copy since it'll get overwritten anyway
            copy.append(None)
        else:
            copy.append(v)
    ret = []
    # we cartesian product all the list arguments
    # the * is argument destructuring
    for items in itertools.product(*[lst[i] for i in lindices]):
        # we make a copy of the list
        curcopy = copy.copy()
        for i, item in zip(lindices, items):
            # we assign the elements
            curcopy[i] = item
        # we append the copy to the return
        ret.append(curcopy)
    return ret

test = ['A', 'B', ['a0', 'a1', 'a2'], 'C', 'D', ['b0', 'b1', 'b2'], 'E', 'F', ['c0', 'c1', 'c2']]
print(combos(test))

如果您确实希望将返回值中的'a0'包装在一个数组中，那么您可以将ret.append(curcopy)替换为ret.append([curcopy]) 。