如何高效对Pandas DataFrame分组并基于前组计算衍生列？

原始数据 #

import pandas as pd

# 构造初始DataFrame
df = pd.DataFrame({
    'Date': ['2022-07-14 02:34:20.348', '2022-07-14 02:34:20.348',
             '2022-07-14 02:37:20.348', '2022-07-14 02:37:20.348',
             '2022-07-14 02:44:20.348', '2022-07-14 02:44:20.348',
             '2022-07-14 02:54:20.348', '2022-07-14 02:54:20.348'],
    'id1': ['A', 'B', 'A', 'B', 'A', 'B', 'A', 'B'],
    'amount': [87.11, 77.12, 89.11, 87.12, 81.11, 87.12, 99.11, 99.11],
    'is_winner': [False, True, False, True, True, False, False, True]
})

# 转换Date为datetime类型（确保时间计算准确）
df['Date'] = pd.to_datetime(df['Date'])

需求说明 #

数据已按Date排序，需生成以下衍生列：

• delta_amount_from_last_win：当前amount与该id1上一次获胜时amount的差值，首次记录为0
• num_wins_before：该id1在当前分组之前的获胜总次数
• num_groups_before：当前分组之前的历史分组总数（按Date分组的序号）
• rank_in_prev_group：该id1在上一组中的获胜排名（0为胜者，1为败者），首次记录为NaN
• time_from_prev(minutes)：当前分组与上一组的时间间隔（分钟），首次记录为NaN

高效实现代码 #

# 1. 给每个Date分组分配唯一序号，方便组间关联
df['group_num'] = df.groupby('Date').ngroup()

# 按id1分组处理个体历史数据
id_group = df.groupby('id1')

# 生成num_wins_before：当前组之前的获胜次数，用cumsum+shift实现矢量化计算
df['num_wins_before'] = id_group['is_winner'].cumsum().shift(1).fillna(0).astype(int)

# 生成delta_amount_from_last_win：当前amount与上一次获胜amount的差值
# 先提取每个id获胜时的amount，向前填充得到最近一次获胜的amount
last_win_amount = id_group.apply(lambda x: x['amount'].where(x['is_winner']).ffill()).reset_index(drop=True)
df['delta_amount_from_last_win'] = df['amount'] - last_win_amount.shift(1).fillna(df['amount'])
# 首次分组记录强制设为0
df.loc[df['group_num'] == 0, 'delta_amount_from_last_win'] = 0

# 生成num_groups_before：直接用分组序号即可
df['num_groups_before'] = df['group_num']

# 生成rank_in_prev_group：上一组的获胜排名，shift整组数据实现组间传递
prev_group_winner = df.groupby('group_num')['is_winner'].transform(lambda x: x).shift(len(df['id1'].unique()))
df['rank_in_prev_group'] = (~prev_group_winner).astype(int).where(df['group_num'] > 0, pd.NA)

# 生成time_from_prev(minutes)：计算组间时间差并映射到每行
group_start_times = df.groupby('group_num')['Date'].min()
group_time_diff = group_start_times.diff().dt.total_seconds() / 60
df['time_from_prev(minutes)'] = df['group_num'].map(group_time_diff).where(df['group_num'] > 0, pd.NA)

# 清理临时列，调整列顺序匹配需求输出
df = df.drop('group_num', axis=1)
target_cols = [
    'Date', 'id1', 'amount', 'is_winner',
    'delta_amount_from_last_win', 'num_wins_before',
    'num_groups_before', 'rank_in_prev_group',
    'time_from_prev(minutes)'
]
df = df[target_cols]

# 查看结果
print(df)

关键优化点 #

• 全程使用Pandas矢量化函数（cumsum、shift、ffill等），避免逐行循环，大幅提升处理效率
• 用ngroup()生成分组序号，简化组间关联逻辑
• 组间数据传递通过shift整组数据实现，避免复杂的分组合并操作

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