使用 pandas groupby + apply 和 condensing groups 计算平均值的更快方法

Posted 2023-03-11

【中文标题】使用 pandas groupby + apply 和 condensing groups 计算平均值的更快方法

【英文标题】：Faster way of computing the mean with pandas groupy + apply and condensing groups

【发布时间】：2021-02-02 05:27:23

【问题描述】：

我想对两个值进行分组，如果该组包含多个元素，则仅返回该组的第一行，并将该值替换为该组的平均值。如果只有一个元素，我想直接返回。我的代码如下所示：

final = df.groupby(["a", "b"]).apply(condense).drop(['a', 'b'], axis=1).reset_index()

def condense(df):
    if df.shape[0] > 1:
        mean = df["c"].mean()
        record = df.iloc[[0]]
        record["c"] = mean
        return(record)
    else:
        return(df)

df 看起来像这样：

a      b     c   d
"f"   "e"    2   True
"f"   "e"    3   False
"c"   "a"    1   True

由于数据框很大，我有73800个组，整个groupby + apply的计算大约需要一分钟。这太长了。有没有办法让它跑得更快？

【参考方案1】：

我认为一个值的平均值与多个值的平均值相同，因此您可以通过 GroupBy.agg 和 mean 为列 c 简化解决方案，而所有其他值由 first 聚合：

d = dict.fromkeys(df.columns.difference(['a','b']), 'first')
d['c'] = 'mean'
print (d)
'c': 'mean', 'd': 'first'

df = df.groupby(["a", "b"], as_index=False).agg(d)
print (df)
   a  b    c     d
0  c  a  1.0  True
1  f  e  2.5  True

【讨论】：

哇，从 58 秒缩短到 0.1 秒，谢谢！