如何在Pandas中分组并在组内过滤后添加聚合列

给定如下DataFrame：

import pandas as pd

df = pd.DataFrame({
    "col1": [1, 2, 3, 4, 5],
    "col2": ["b", "b", "b", "a", "b"],
    "col3": [6, 2, 11, 1, 3]
})

需要新增一列mean_col，规则为：对每一行，计算同col2分组下、col3值小于当前行col3值的col1的均值。预期结果如下：

pd.DataFrame({
    "col1": [1, 2, 3, 4, 5],
    "col2": ["b", "b", "b", "a", "b"],
    "col3": [6, 2, 11, 1, 3],
    "mean_col": [3.5, None, 2.666, None, 2.0]
})

已通过iterrows循环实现，但希望使用更高效的向量化/分组操作替代遍历。

解决方案1：分组后结合apply处理单组逻辑 #

利用groupby('col2')分组后，对每组内的行单独筛选符合条件的记录并计算均值：

def get_group_mean(row, group_data):
    # 筛选组内col3小于当前行col3的col1值
    valid_col1 = group_data.loc[group_data['col3'] < row['col3'], 'col1']
    return valid_col1.mean() if not valid_col1.empty else None

# 分组后逐行应用计算逻辑
df['mean_col'] = df.groupby('col2').apply(
    lambda g: g.apply(lambda row: get_group_mean(row, g), axis=1)
).reset_index(level=0, drop=True)

# 保留三位小数对齐预期结果
df['mean_col'] = df['mean_col'].round(3)
print(df)

输出结果：

col1 col2  col3  mean_col
0     1    b     6      3.500
1     2    b     2        NaN
2     3    b    11      2.667
3     4    a     1        NaN
4     5    b     3      2.000

解决方案2：自连接+分组聚合（高性能向量化方案） #

通过自连接生成同组内的所有行对，筛选符合条件的记录后批量计算均值，完全避免行级遍历：

# 自连接，保留同col2的所有行组合
merged_df = df.merge(df, on='col2', suffixes=('', '_other'))
# 筛选出另一行col3小于当前行的记录
filtered = merged_df[merged_df['col3_other'] < merged_df['col3']]
# 按原行索引分组，计算符合条件的col1均值
mean_result = filtered.groupby(filtered.index)['col1_other'].mean().round(3)
# 将结果合并回原DataFrame，无匹配项填充NaN
df['mean_col'] = mean_result
print(df)

该方法依赖pandas的向量化操作，在数据量较大时性能远优于行遍历。

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