如何修改自定义函数适配pandas groupby实现分组独立计算

核心问题原因 #

之前分组计算失败主要是三个原因：

• 原函数使用链式索引（df['Units'][i]）取值赋值，分组切片得到的子表保留原表全局索引，按位置i取数时会出现标签错位，取到错误行的值
• 原函数直接在传入对象上修改，遍历分组时没有做独立副本和索引重置，会出现跨组值继承、赋值覆盖的问题
• 自行遍历分组时没有做索引对齐，计算结果无法正确映射回原表对应行

可直接运行的实现代码 #

import pandas as pd

# 构造示例数据
d = {'type': ['A', 'A', 'A' ,'A' , 'A', 'A', 'A','A' ,'A' ,'A', 'A', 'A', 'A','B', 'B', 'B' ,'B' , 'B', 'B', 'B','B' ,'B' ,'B', 'B', 'B', 'B'], 
 'Date': ['Jun-21','Jul-21','Aug-21','Sep-21','Oct-21','Nov-21','Dec-21','Jan-22','Feb-22','Mar-22','Apr-22','May-22','Jun-22', 'Jun-21','Jul-21','Aug-21','Sep-21','Oct-21','Nov-21','Dec-21','Jan-22','Feb-22','Mar-22','Apr-22','May-22','Jun-22'],
 'Units':[0, 0, 0, 0, 10, 0, 20, 0, 0, 7, 12, 35, 0, 0,0,0,0,7,4,0,4,9,5,8,3,11]}
df = pd.DataFrame(data=d)

# 适配分组的计算函数
def calc_group_expected(group):
    # 重置索引避免位置取值错位，创建副本避免修改原分组对象
    g = group.reset_index(drop=True).copy()
    g['Expected'] = 0
    for i in range(1, len(g)):
        prev_exp = g.loc[i-1, 'Expected']
        curr_units = g.loc[i, 'Units']
        if curr_units == 0:
            g.loc[i, 'Expected'] = prev_exp
        if curr_units > 0:
            # 原逻辑中对单个标量调用cumsum是冗余操作，直接取max做下限截断即可，和单组计算结果完全一致
            add = max((curr_units - 2 * prev_exp) // 5, 0)
            g.loc[i, 'Expected'] = prev_exp + add
    # 恢复原分组索引，保证结果能正确对齐原表
    g.index = group.index
    return g

# 按type分组逐组应用计算，不保留分组键索引
df = df.groupby('type', group_keys=False).apply(calc_group_expected)

结果说明 #

运行上述代码后得到的结果和给出的预期输出完全匹配，每个分组独立初始化Expected=0，递推计算不会跨组串值。
如果数据量特别大，可以后续把逐行递推逻辑优化成numba加速的累计计算，但当前groupby+apply的写法和原有单组计算逻辑完全对齐，可读性最高，调试成本最低，是普通数据量下的最优选择。

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