高效实现Pandas列按年度分组占比计算的方法

问题场景 #

手里有个带year、cc_rating、number_x字段的Pandas DataFrame，要算每条记录的number_x在对应年度的占比。因为year不唯一没法直接设索引做除法，之前用merge加groupby的方法跑起来效率太低，想找更优的实现方式。

示例数据：

df = pd.DataFrame({
    "year": {"0":2005,"1":2005,"2":2005,"3":2006,"4":2006,"5":2006,"6":2007,"7":2007,"8":2007},
    "cc_rating": {"0":"2","1":"2a","2":"2b","3":"2","4":"2a","5":"2b","6":"2","7":"2a","8":"2b"},
    "number_x": {"0":9368,"1":21643,"2":107577,"3":10069,"4":21486,"5":110326,"6":10834,"7":21566,"8":111082}
})

原低效代码 #

df= pd.merge(df, df.groupby('year').sum(), left_on='year',right_index=True)
df['%'] = round((df['number_x'] / df['number_y'])*100 , 2)
df = df.drop('number_y', axis=1)

优化方案：用groupby.transform替代merge #

原方法的问题在于merge会生成临时表，数据量大时开销很高。而transform方法可以直接把分组求和的结果广播匹配到原表的对应行，全程不需要额外的合并操作，效率提升明显。

优化后的代码：

# 计算每条记录对应年度的number_x总和，自动匹配到对应行
year_total = df.groupby('year')['number_x'].transform('sum')
# 直接计算占比
df['%'] = round((df['number_x'] / year_total) * 100, 2)

效果说明 #

这个方法和原代码得到的结果完全一致，但省去了merge和删除临时列的步骤，代码更简洁，大数据量下的运行速度会快很多——尤其是当你的DataFrame有几万甚至几十万行时，差距会非常明显。

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