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

Question

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

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的计算大约需要一分钟。这太长了。有没有办法让它运行得更快？

Answer 1

我认为一个值的平均值与多个值的平均值相同，因此您可以通过 GroupBy.agg 来简化解决方案列 c 的 mean 和所有其他值按 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