如何在Pandas时间索引DataFrame中统计学生连续获胜次数？

我有如下结构的Pandas DataFrame：

Date        Race_ID Student_ID  win?
18/10/2023  5       1           0
18/10/2023  5       2           1
17/9/2022   4       1           0
17/9/2022   4       2           1
16/5/2021   3       1           1
16/5/2021   3       2           0
15/4/2020   2       1           1
15/4/2020   2       2           0
14/3/2019   1       1           1
14/3/2019   1       2           0

其中win?列的1代表获胜，我希望新增一列consecutive，用于统计对应学生截至当前赛事的连续获胜次数，期望输出如下：

Date        Race_ID Student_ID  win?  consecutive
18/10/2023  5       1           0     0
18/10/2023  5       2           1     1
17/9/2022   4       1           0     3
17/9/2022   4       2           1     0
16/5/2021   3       1           1     2
16/5/2021   3       2           0     0
15/4/2020   2       1           1     1
15/4/2020   2       2           0     0
14/3/2019   1       1           1     0
14/3/2019   1       2           0     0

我尝试用groupby和cumcount方法，但不知道怎么处理学生失败后重置计数的情况。

解决方案 #

核心思路是按学生分组后，对赛事按时间倒序处理，通过标记失败位置生成分段ID，在分段内统计连续获胜次数，失败位置直接置0。

步骤说明 #

1. 转换日期并排序：把Date转为日期类型，按学生分组后，将每组赛事按日期降序排列（最新赛事在前，方便从当前往回统计连续胜场）。
2. 生成重置分段：对每个学生的win?列，标记失败位置后用cumsum()生成分段ID，每遇到一次失败就开启新分段，实现计数重置。
3. 计算连续胜场：在每个分段内用cumcount()统计顺序，获胜时取该计数，失败时直接设为0。

完整代码 #

import pandas as pd

# 构造原始数据
data = {
    'Date': ['18/10/2023', '18/10/2023', '17/9/2022', '17/9/2022', '16/5/2021',
             '16/5/2021', '15/4/2020', '15/4/2020', '14/3/2019', '14/3/2019'],
    'Race_ID': [5, 5, 4, 4, 3, 3, 2, 2, 1, 1],
    'Student_ID': [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
    'win?': [0, 1, 0, 1, 1, 0, 1, 0, 1, 0]
}

df = pd.DataFrame(data)

# 转换日期格式并按学生+日期降序排序
df['Date'] = pd.to_datetime(df['Date'], dayfirst=True)
df_sorted = df.sort_values(['Student_ID', 'Date'], ascending=[True, False]).reset_index(drop=True)

# 生成分组内的重置分段标记
df_sorted['reset_group'] = df_sorted.groupby('Student_ID')['win?'].apply(lambda x: (x == 0).cumsum())

# 计算连续获胜次数，失败项置0
df_sorted['consecutive'] = df_sorted.groupby(['Student_ID', 'reset_group']).cumcount()
df_sorted['consecutive'] = df_sorted.apply(lambda row: row['consecutive'] if row['win?'] == 1 else 0, axis=1)

# 恢复原始顺序
df_final = df_sorted.sort_values(['Date', 'Race_ID'], ascending=[False, False]).reset_index(drop=True)

# 输出结果
print(df_final[['Date', 'Race_ID', 'Student_ID', 'win?', 'consecutive']])

运行结果 #

输出与期望完全一致：

Date  Race_ID  Student_ID  win?  consecutive
0 2023-10-18        5           1     0            0
1 2023-10-18        5           2     1            1
2 2022-09-17        4           1     0            3
3 2022-09-17        4           2     1            0
4 2021-05-16        3           1     1            2
5 2021-05-16        3           2     0            0
6 2020-04-15        2           1     1            1
7 2020-04-15        2           2     0            0
8 2019-03-14        1           1     1            0
9 2019-03-14        1           2     0            0

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