如何在Pandas中按条件分组并应用自定义聚合函数

给定如下定义的DataFrame：

import pandas as pd

df = pd.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar', 'foo', 'bar','foo', 'bar'],
                   'B' : [1, 2, 3, 4, 5, 6, 7, 8],
                   'C' : ['mean', 'halfmean', 'mean', 'halfmean', 'mean', 'halfmean', 'mean', 'halfmean']})

需求：按A列分组后，根据C列的条件聚合B列数据：

• foo分组：计算B列所有值的均值
• bar分组：计算B列最大半数值的均值（注：同一分组的C列值完全一致）

现有分组代码框架：

grouped = df.groupby('A')
????

预期结果：

foo 4 #(1 + 3 + 5 + 7)/4
bar 7 #(6 + 8)/2

方法一：分组映射聚合（性能更优） #

先提取每个分组对应的聚合类型，通过字典映射调用对应逻辑，利用Pandas内置矢量化方法提升效率：

# 1. 获取每个A分组对应的C列聚合类型
agg_type_map = df.groupby('A')['C'].first().to_dict()

# 2. 定义不同聚合类型的处理逻辑
agg_logic = {
    'mean': lambda col: col.mean(),
    'halfmean': lambda col: col.nlargest(len(col)//2).mean()
}

# 3. 执行分组聚合
result = df.groupby('A')['B'].agg(lambda x: agg_logic[agg_type_map[x.name]](x))

方法二：自定义分组处理函数（逻辑直观） #

通过apply传入自定义函数，直接根据分组内的C值处理B列：

def group_agg(group):
    agg_type = group['C'].iloc[0]
    if agg_type == 'mean':
        return group['B'].mean()
    elif agg_type == 'halfmean':
        # 取排序后后半段数值的均值
        return group['B'].sort_values().iloc[len(group)//2:].mean()

result = df.groupby('A').apply(group_agg)

结果验证 #

运行代码后输出结果：

print(result)
# 输出：
# A
# bar    7.0
# foo    4.0
# dtype: float64

性能说明 #

方法一效率更高，避免了apply的逐组循环开销，利用nlargest等内置矢量化方法处理数据，大数据量场景下优势明显；方法二逻辑更直观，适合快速实现和调试。

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