如何高效计算Pandas DataFrame中分组聚合列的占比？

原问题场景 #

你有如下Pandas DataFrame，需要计算share列（每组内各amount除以该组end_amount对应的值），当前用循环实现但效率极低：

import pandas as pd

d = {"col1":["A", "A", "A", "B", "B", "B"], "col2":["start_amount", "mid_amount", "end_amount", "start_amount", "mid_amount", "end_amount"], "amount":[0, 2, 8, 1, 2, 3]}
df_test = pd.DataFrame(d)

df_test["share"] = 0
for i in range(len(df_test)):
    df_test.loc[i, "share"] = df_test.loc[i, "amount"] / df_test.loc[(df_test["col1"] == df_test.loc[i, "col1"]) & (df_test["col2"] == "end_amount"), "amount"].values

优化方案 #

利用Pandas的矢量化操作和分组映射替代循环，大幅提升效率：

1. 先提取每个col1分组对应的end_amount值，构建映射关系：

# 筛选出end_amount行，以col1为索引，保留amount列
end_mapping = df_test[df_test['col2'] == 'end_amount'].set_index('col1')['amount']

2. 直接通过映射将每个行的col1对应到end_amount，再做除法计算share：

df_test['share'] = df_test['amount'] / df_test['col1'].map(end_mapping)

结果验证 #

执行后df_test的share列结果和原循环完全一致，但效率提升显著——尤其是当数据量达到万级以上时，矢量化操作的速度会是循环的几十甚至上百倍。

原理说明 #

Pandas的核心优势是基于Numpy的矢量化运算，避免了Python层面的逐行循环开销。通过map方法可以高效地将分组的基准值（end_amount）匹配到每一行，再进行批量除法运算，完全利用了底层的C语言优化。

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