Pandas按指定列分组后仅对组内首行执行apply并同步结果的优化方法

实现方案 #

你可以通过「先批量计算唯一分组的结果，再回填到原表」的逻辑实现，仅需要对分组级别的少量数据执行apply，大幅降低计算量，完整可运行代码如下：

import pandas as pd
from io import StringIO

# 原有DataFrame构造逻辑
df_str = """
  ValOption  RB test contrat
0       SLA  4  3    23
1       AC   5  4    12
2       SLA  5  5    23
3       AC   2  4    39
4       SLA  5  5    26
5       AC   5  4    52
6       SLA  4  3    64
"""
df = pd.read_csv(StringIO(df_str.strip()), sep='\s+')

# 原有自定义函数无需修改
def func(row):
    v1 = row['RB'] * 3
    v2 = row['test'] - 1
    return v1 + v2

# ----------------------优化实现逻辑----------------------
# 1. 定义分组键
group_cols = ['ValOption', 'RB']
# 2. 提取所有唯一分组的首行，仅保留自定义函数需要用到的列，可根据实际需求调整保留列
unique_groups = df.groupby(group_cols, as_index=False).first()[group_cols + ['test']]
# 3. 仅对少量分组首行执行apply计算，得到每个分组对应的结果
unique_groups['group_result'] = unique_groups.apply(lambda row: func(row), axis=1)
# 4. 将分组结果关联回填到原DataFrame的所有行
df = df.merge(unique_groups[group_cols + ['group_result']], on=group_cols, how='left')
# 可选：重命名为目标列名
df = df.rename(columns={'group_result': 'new_col'})

补充说明 #

运行后得到的new_col和全量apply得到的结果完全一致，计算量从「原表行数次apply」降低为「分组数个apply」，数据量越大、单分组行数越多，性能提升越明显。
如果需要保留原表的排序，可在merge后按原索引重排：

df = df.sort_index()

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