Pandas groupby聚合列表列如何实现逐元素求和而非拼接

问题原因 #

Python中列表的默认加法操作为拼接，因此pandas原生sum聚合会将分组内所有列表直接拼接，无法实现逐元素求和。

解决方法 #

自定义聚合函数实现逐位求和，有两种常用实现方式：

方法1：依赖numpy实现（运算效率更高）

import pandas as pd
import numpy as np

# 原有DataFrame
df  = pd.DataFrame({'upc':[1,1,1],'store':[1,2,3],'date':['jan','jan','jan'],'pred':[[1,1,1],[2,2,2],[3,3,3]],'act':[[4,4,4],[5,5,5],[6,6,6]]})

# 分组聚合
res = df.groupby(by = ["upc","date"]).agg(
    pred = ("pred", lambda x: np.array(x.tolist()).sum(axis=0).tolist()),
    act = ("act", lambda x: np.array(x.tolist()).sum(axis=0).tolist())
).reset_index()

方法2：纯Python实现（无需额外依赖）

res = df.groupby(by = ["upc","date"]).agg(
    pred = ("pred", lambda x: [sum(i) for i in zip(*x)]),
    act = ("act", lambda x: [sum(i) for i in zip(*x)])
).reset_index()

输出结果 #

打印res即可得到预期效果：

upc date       pred           act
0    1  jan  [6, 6, 6]  [15, 15, 15]

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