如何向量化Pandas DataFrame多列聚合操作以提升计算效率

原始数据与需求 #

给定以下以布尔列为主的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({"A": [1, 2, 3, 1, 2, 3],
                   "B": ['a', 'b', 'c', 'a', 'b', 'c'],
                   "f1": [True, True, True, True, True, False],
                   "f2": [True, True, True, True, False, True],
                   "f3": [True, True, True, False, True, True],
                   "f4": [True, True, False, True, True, True],
                   "f5": [True, False, True, True, True, True],
                   "target1": [True, False, True, True, False, True],
                   "target2": [False, True, True, False, True, False]})

需求是：针对每个f1-f5列的True/False类别，按(A,B)分组后对target1、target2列求和。原实现通过遍历f列完成，现需替换为向量化操作以减少耗时。

原始循环实现 #

for col in ["f1", "f2", "f3", "f4", "f5"]:
    print(col, "\n", 
          df[df[col]].groupby(["A", "B"]).agg({"target1": "sum", "target2": "sum"}), "\n",
          df[~df[col]].groupby(["A", "B"]).agg({"target1": "sum", "target2": "sum"}))

向量化优化方案 #

利用Pandas的melt函数将宽格式的f列转换为长格式，再通过一次分组聚合完成所有计算，完全避免循环：

# 1. 将f列转换为长格式，保留分组列和目标列
melted = df.melt(
    id_vars=['A', 'B', 'target1', 'target2'],
    value_vars=['f1', 'f2', 'f3', 'f4', 'f5'],
    var_name='f_col',  # 存储原f列的列名
    value_name='f_flag'  # 存储原f列的布尔值（True/False）
)

# 2. 一次性分组聚合所有情况
result = melted.groupby(['A', 'B', 'f_col', 'f_flag']).agg(
    target1_sum=('target1', 'sum'),
    target2_sum=('target2', 'sum')
).reset_index()

# 查看结果
print(result)

结果说明 #

执行后会得到包含所有f列、所有布尔分组下的聚合结果，示例片段如下：

A  B f_col  f_flag  target1_sum  target2_sum
0  1  a    f1    True            2            0
1  1  a    f2    True            2            0
2  1  a    f3    True            1            0
3  1  a    f3   False            1            0
4  1  a    f4    True            2            0
...

如果需要更贴近原循环的输出格式（按f列分别展示True/False结果），可以对结果进行重塑：

pivoted_result = result.pivot(
    index=['A', 'B'],
    columns=['f_col', 'f_flag'],
    values=['target1_sum', 'target2_sum']
)
print(pivoted_result)

优势 #

这种向量化操作将所有f列的计算合并为一次批量处理，耗时显著低于循环实现，接近处理单个f列的时间开销。

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