如何在Pandas中高效添加基于多分组求和的新列value3

需求说明 #

现有Pandas DataFrame包含cat1、cat2、cat3、cat4、value1、value2列，需要添加新列value3，其值为相同cat1、cat2、cat3分组下所有cat4对应的value1之和，要求采用向量化高效实现，避免编写Python循环。

示例输入数据 #

cat1, cat2, cat3, cat4, value1, value2
A,A,A,A, 3,1
A,A,A,B, 5,2
A,A,A,C, 22,4
A,A,X,A, 0,5
A,A,X,B, 4,5
A,A,X,C, 6,5

期望输出结果 #

cat1, cat2, cat3, cat4, value1, value2, value3
A,A,A,A, 3,1, 30
A,A,A,B, 5,2, 30
A,A,A,C, 22,4, 30
A,A,X,A, 0,5, 10
A,A,X,B, 4,5, 10
A,A,X,C, 6,5, 10

解决方案 #

使用Pandas内置的groupby()结合transform()方法即可实现，这是完全向量化的操作，底层基于C实现，性能远优于Python循环：

import pandas as pd

# 读取示例数据（也可直接使用已有DataFrame）
df = pd.read_csv(pd.io.common.StringIO("""
cat1, cat2, cat3, cat4, value1, value2
A,A,A,A, 3,1
A,A,A,B, 5,2
A,A,A,C, 22,4
A,A,X,A, 0,5
A,A,X,B, 4,5
A,A,X,C, 6,5
"""), skipinitialspace=True)

# 计算分组求和并添加value3列
df['value3'] = df.groupby(['cat1', 'cat2', 'cat3'])['value1'].transform('sum')

# 查看结果
print(df.to_csv(index=False))

原理说明 #

• groupby(['cat1', 'cat2', 'cat3'])：按照指定的三个分类列完成分组
• ['value1'].transform('sum')：对每个分组的value1列求和，随后将求和结果自动广播匹配到原DataFrame的对应分组行，生成与原表长度一致的Series，直接赋值即可得到目标列。

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