如何使用itertools实现多列表各选指定数量元素的组合生成

有简洁的实现方式，无需手动拆分列表循环调用接口，通过itertools内置接口组合即可完成需求：

实现代码 #

import itertools

# 输入参数
a = [[1,2,3,4,5],[6,7,8,9,10]]
pick_nums = [2, 3] # 顺序对应每个子列表需要选取的元素数量

# 生成最终结果
r = [
    sum(comb_group, []) 
    for comb_group in itertools.product(
        *[itertools.combinations(sublist, n) for sublist, n in zip(a, pick_nums)]
    )
]

逻辑说明 #

1. 内层列表推导先遍历每个子列表和对应选取数量，调用itertools.combinations生成该子列表所有指定长度的组合集合
2. 调用itertools.product对所有子列表的组合集合求笛卡尔积，得到的每个comb_group格式为(子列表1的选中组合, 子列表2的选中组合, ...)
3. 用sum(comb_group, [])将多组组合拼接为单个列表，若处理的子列表/元素数量较大，可将拼接逻辑替换为list(itertools.chain.from_iterable(comb_group))获得更高性能

上述实现支持任意长度的嵌套列表，只需保证pick_nums的长度和a的子列表数量一致即可。

内容的提问来源于stack exchange，提问作者Marcelo Ruiz