在Jupyter Notebook中用Python生成问卷全答案组合遇内存错误及筛选需求

问题背景 #

要生成含约40个问题的问卷所有答案组合（部分问题3个选项，部分约10个，答案含字符串和正自然数），输出到Excel时原代码触发内存错误，同时需要添加条件筛选规则（如Question 0为"No"或"N/A"时，Question 1和Question 2必须为0）。

原错误代码：

import pandas as pd
from itertools import product
df = pd.read_excel("Test2023 PYTHON.xlsx")
uniques = [df[i].unique().tolist() for i in df.columns]
combinations = pd.DataFrame(product(*uniques), columns = df.columns)
combinations_df.to_excel("combinations.xlsx", index=False)

核心原因 #

40个问题的组合数是天文数字（即使每个问题平均5个选项，总组合数达5^40≈9e27），一次性将所有组合加载到DataFrame会直接耗尽内存。

解决方案：迭代生成+实时筛选+分批写入 #

通过逐一生成组合、实时筛选符合条件的结果、分批写入Excel的方式，避免内存过载。

完整代码实现 #

import pandas as pd
from itertools import product
from openpyxl import load_workbook

# 读取问卷选项数据
df = pd.read_excel("Test2023 PYTHON.xlsx")
uniques = [col.unique().tolist() for col in df.itertuples(index=False, name=None)]
cols = df.columns.tolist()

# 定义条件筛选规则
def meets_requirements(combo):
    combo_dict = dict(zip(cols, combo))
    # 规则1：Question 0为"No"或"N/A"时，Question1和Question2必须为0
    if combo_dict["Question 0"] in ["No", "N/A"]:
        if combo_dict["Question 1"] != 0 or combo_dict["Question 2"] != 0:
            return False
    # 可在此添加其他筛选规则，示例：
    # if combo_dict["Question 5"] == "Yes":
    #     if combo_dict["Question 6"] not in ["A", "B"]:
    #         return False
    return True

# 配置分批写入参数
batch_size = 10000  # 每批写入的行数，可根据内存调整
output_path = "filtered_combinations.xlsx"
current_batch = []

# 初始化空Excel文件并写入表头
pd.DataFrame(columns=cols).to_excel(output_path, index=False)

# 遍历所有组合，筛选后分批写入
for combo in product(*uniques):
    if meets_requirements(combo):
        current_batch.append(combo)
        # 批次满额时写入Excel
        if len(current_batch) >= batch_size:
            book = load_workbook(output_path)
            with pd.ExcelWriter(output_path, engine="openpyxl", mode="a", if_sheet_exists="overlay") as writer:
                pd.DataFrame(current_batch, columns=cols).to_excel(writer, index=False, header=False, startrow=writer.sheets["Sheet1"].max_row)
            current_batch = []

# 写入剩余未达批次的结果
if current_batch:
    book = load_workbook(output_path)
    with pd.ExcelWriter(output_path, engine="openpyxl", mode="a", if_sheet_exists="overlay") as writer:
        pd.DataFrame(current_batch, columns=cols).to_excel(writer, index=False, header=False, startrow=writer.sheets["Sheet1"].max_row)

代码说明 #

1. 迭代生成组合：用itertools.product逐个生成组合，而非一次性生成全部，大幅降低内存占用。
2. 实时条件筛选：通过自定义函数meets_requirements过滤不符合规则的组合，减少后续写入的数据量。
3. 分批写入Excel：每积累一定数量的有效组合就写入文件，避免内存中堆积过多数据。

注意事项 #

• 可根据机器内存大小调整batch_size，内存充足时可适当调大以提升效率。
• 筛选规则需根据实际需求补充，确保逻辑覆盖所有业务要求。
• 若部分问题选项极多，可先对这些问题做预筛选，进一步减少迭代的总组合数。

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