Pandas按行多列运算实现分组加权浓度的简便方法

问题复现 #

测试DataFrame构造代码如下：

import pandas as pd
df = pd.DataFrame(list(zip(['a', 'a', 'b', 'b', 'c', 'c', 'c'], 
                           ['a1', 'a2', 'b1', 'b2', 'c1', 'c2', 'c3'],
                           [110, 80, 100, 180, 12], 
                           [5, 7, 2, 6, 10])), 
                      columns=['name', 'ingredient', 'amount', 'con'])

计算需求：按name字段分组，基于组内各原料的amount（含量）和con（浓度）计算分组总浓度。原有分步实现逻辑为逐行计算amount*con、关联组内总含量、计算单原料浓度占比后分组求和，得到如下输出：

name
a     5.842105
b     4.571429
c    10.000000
Name: what_i_want, dtype: float64

简洁实现方案 #

你当前的分步计算本质是求以amount为权重的con加权平均值，数学上可化简为「组内amount*con的总和 / 组内amount总和」，完全不需要merge操作，以下两种写法均可得到完全一致的结果：

写法1：单行聚合（代码最短） #

直接在分组聚合时同步计算分子、分母，无需生成任何中间列：

res = df.groupby('name').apply(lambda x: (x['amount'] * x['con']).sum() / x['amount'].sum())

写法2：transform实现（大数据量性能更好） #

如果需要保留逐行计算的中间过程，用groupby.transform直接生成组内总含量列即可，比merge执行效率更高，数据量越大性能差距越明显：

df['amt_sum'] = df.groupby('name')['amount'].transform('sum')
res = (df['amount'] * df['con'] / df['amt_sum']).groupby(df['name']).sum()

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