如何使用itertools.product生成符合要求的Pandas DataFrame？

问题分析 #

当前用itertools.product(arr1, arr2)得到的是笛卡尔积，会把arr1的每个元素和arr2的所有子列表配对，这和预期的「arr1元素与对应位置arr2子列表的逐个元素配对」逻辑不符。

正确的itertools实现方案 #

我们需要针对arr1和arr2的对应元素对，分别生成元素级的配对，再用itertools.chain.from_iterable把所有配对结果拼接起来：

import pandas as pd
from itertools import product, chain

arr1 = ["A","B"]
arr2 = [[1,2],[3,4,5]]

# 对每一组(a1, sub_arr)，生成(a1, num)的配对，再拼接所有结果
data = chain.from_iterable(product([a1], sub_arr) for a1, sub_arr in zip(arr1, arr2))
df2 = pd.DataFrame(data, columns=["short_list", "long_list"])

输出验证 #

执行后得到的df2与预期一致：

逻辑说明 #

• zip(arr1, arr2)：将arr1的元素和arr2对应位置的子列表配对（即("A", [1,2])、("B", [3,4,5])）
• product([a1], sub_arr)：对每一对，生成a1和sub_arr中每个元素的配对（比如第一组生成("A",1)、("A",2)）
• chain.from_iterable()：把所有组的配对结果展开成扁平迭代器，直接传入pd.DataFrame即可

内容的提问来源于stack exchange，提问作者Gооd_Mаn