如何加速Pandas DataFrame双重循环运算？求并行优化方案

你的需求是对每个元素计算(行和 × 列和) / 整体总和，原三重循环因为Python层面的循环开销和重复计算，在处理大型DataFrame时速度极差。以下是两种更高效的优化方案，优先推荐矢量化计算，并行化作为补充场景使用：

方案1：矢量化计算（最优选择） #

利用NumPy和Pandas的矢量化操作，完全避免Python循环，底层由C实现，速度提升几个数量级。核心思路是先一次性计算所有行和、列和，再通过外积生成整个期望矩阵：

import numpy as np
import pandas as pd

# 先初始化结果结构，与原数据结构匹配
df_aug15_exp = {}

for h in header:
    df_sub = df_aug15[h]
    # 计算整体总和（仅需一次）
    total = df_sub.sum().sum()
    # 计算所有行的和（形状：(行数,)）
    row_sums = df_sub.sum(axis=1)
    # 计算所有列的和（形状：(列数,)）
    col_sums = df_sub.sum(axis=0)
    # 行和与列和做外积，得到期望矩阵，再除以总和
    exp_matrix = np.outer(row_sums, col_sums) / total
    # 转换为带原索引和列名的DataFrame
    df_aug15_exp[h] = pd.DataFrame(exp_matrix, index=df_sub.index, columns=df_sub.columns)

为什么比原代码快？ #

• 原代码每个元素都重复计算行和、列和、总和，这里仅计算一次行和、列和、总和，避免大量重复运算
• 外积运算np.outer是底层矢量化操作，比Python循环快数十到数百倍

方案2：并行化处理（适用于多子DataFrame场景） #

如果header数量极多，且每个子DataFramedf_aug15[h]都很大，可以用多进程并行处理每个子任务，抵消进程开销后能进一步提升速度：

import concurrent.futures
import numpy as np
import pandas as pd

def compute_expected_sub(df_sub):
    total = df_sub.sum().sum()
    row_sums = df_sub.sum(axis=1)
    col_sums = df_sub.sum(axis=0)
    exp_matrix = np.outer(row_sums, col_sums) / total
    return pd.DataFrame(exp_matrix, index=df_sub.index, columns=df_sub.columns)

# 使用进程池并行处理每个header对应的子DataFrame
with concurrent.futures.ProcessPoolExecutor() as executor:
    # 提交所有任务
    task_map = {h: executor.submit(compute_expected_sub, df_aug15[h]) for h in header}
    # 收集结果
    df_aug15_exp = {h: task_map[h].result() for h in header}

注意事项 #

• 并行化适合header数量多的场景，如果只有少量子DataFrame，进程启动和数据传递的开销会抵消并行收益
• 若使用Jupyter Notebook，并行化需要注意环境兼容性，可能需要调整启动方式

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