Pandas按Name分组计算带Flag条件的年龄差值问题

问题背景 #

给定如下Pandas DataFrame：

需求：为每行计算当前行该人员的Age与同组（Name相同）之前行中最后一个Flag为1的Age的差值；若该人员之前无Flag为1的行，结果为0。期望输出结果如下：

尝试使用groupby+expanding.apply时遇到问题：传给自定义函数的参数仅为Series，无法同时访问Flag和Age列来完成计算。

解决方案 #

方法一：组内自定义函数处理（适合小数据集） #

直接对分组后的每个组应用自定义函数，函数接收整个组的DataFrame，可同时访问所有列：

import pandas as pd

def compute_last_flag_diff(group):
    result = [0]
    # 初始化最后一个Flag=1的Age
    last_flag1_age = group.iloc[0]['Age'] if group.iloc[0]['Flag'] == 1 else None
    
    for idx in range(1, len(group)):
        current_row = group.iloc[idx]
        # 计算差值，无前置值则为0
        diff = current_row['Age'] - last_flag1_age if last_flag1_age is not None else 0
        result.append(diff)
        
        # 更新最后一个Flag=1的Age（当前行Flag=1时）
        if current_row['Flag'] == 1:
            last_flag1_age = current_row['Age']
    
    return pd.Series(result, index=group.index)

# 应用函数到每个分组
df['Result'] = df.groupby('Name', group_keys=False).apply(compute_last_flag_diff)

执行后df['Result']即可得到期望输出。

方法二：向量化操作（高效，适合大数据集） #

利用Pandas的向量化方法替代循环，性能更优：

import pandas as pd

# 1. 仅保留Flag=1的行的Age，其余设为NaN
df['flag1_age'] = df['Age'].where(df['Flag'] == 1)

# 2. 组内向前填充（ffill），得到每行之前最近的Flag=1的Age（注意shift()排除当前行）
df['last_flag1_age'] = df.groupby('Name')['flag1_age'].shift().ffill()

# 3. 计算差值，无前置值的填充为0并转整数
df['Result'] = (df['Age'] - df['last_flag1_age']).fillna(0).astype(int)

# 清理临时列
df.drop(['flag1_age', 'last_flag1_age'], axis=1, inplace=True)

这种方法通过where、shift、ffill等内置方法完成计算，避免了显式循环，处理大数据集时效率更高。

原方法问题说明 #

groupby.expanding.apply默认会将窗口内的每一列单独作为Series传入函数，无法同时获取同一行的Flag和Age值。而直接对分组后的组使用apply，函数会接收整个组的DataFrame，从而可以自由访问所有列的数据。

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