在Python DataFrame中计算网球选手的连胜/连败记录

需求说明 #

现有存储多年网球赛事数据的DataFrame（tennis_data_processed），其中target字段表示player1获胜（1）或失败（0）。需要新增player1_streak和player2_streak字段，规则如下：

• 选手首次出场记为0；
• 连胜则依次递增（1、2…），连败则依次递减（-1、-2…）；
• 胜负转换时重置为对应起始值（胜则1，负则-1）；
• 选手可能出现在player1_id或player2_id列中。

示例数据与预期结果 #

初始数据 #

预期结果 #

代码实现 #

基础版本（适合小数据集） #

import pandas as pd

# 初始化示例数据（替换为你的真实DataFrame）
data = {
    'match_id': [1,2,3,4,5],
    'player1_id': ['A','B','A','A','B'],
    'player2_id': ['B','C','C','B','A'],
    'target': [1,1,0,0,1],
    'match_date': ['2023-01-01','2023-01-02','2023-01-03','2023-01-04','2023-01-05']
}
tennis_data_processed = pd.DataFrame(data)
tennis_data_processed['match_date'] = pd.to_datetime(tennis_data_processed['match_date'])

# 1. 将每场比赛拆分为两个选手的独立胜负记录
player_records = []
for _, row in tennis_data_processed.iterrows():
    # 记录player1的胜负结果
    player_records.append({
        'match_id': row['match_id'],
        'player_id': row['player1_id'],
        'win': 1 if row['target'] == 1 else 0,
        'match_date': row['match_date']
    })
    # 记录player2的胜负结果
    player_records.append({
        'match_id': row['match_id'],
        'player_id': row['player2_id'],
        'win': 0 if row['target'] == 1 else 1,
        'match_date': row['match_date']
    })

player_records_df = pd.DataFrame(player_records)

# 2. 按选手ID和比赛日期排序，确保按时间顺序处理
player_records_df = player_records_df.sort_values(by=['player_id', 'match_date']).reset_index(drop=True)

# 3. 计算每个选手的streak
def calculate_streak(group):
    streaks = [0]  # 首次出场为0
    for i in range(1, len(group)):
        prev_win = group['win'].iloc[i-1]
        curr_win = group['win'].iloc[i]
        prev_streak = streaks[i-1]
        
        if curr_win == 1:
            streaks.append(prev_streak + 1 if prev_win == 1 else 1)
        else:
            streaks.append(prev_streak - 1 if prev_win == 0 else -1)
    group['streak'] = streaks
    return group

player_streaks_df = player_records_df.groupby('player_id', group_keys=False).apply(calculate_streak)

# 4. 将streak合并回原DataFrame
tennis_data_processed['player1_streak'] = tennis_data_processed.apply(
    lambda row: player_streaks_df[(player_streaks_df['match_id'] == row['match_id']) & 
                                 (player_streaks_df['player_id'] == row['player1_id'])]['streak'].values[0],
    axis=1
)
tennis_data_processed['player2_streak'] = tennis_data_processed.apply(
    lambda row: player_streaks_df[(player_streaks_df['match_id'] == row['match_id']) & 
                                 (player_streaks_df['player_id'] == row['player2_id'])]['streak'].values[0],
    axis=1
)

# 查看结果
print(tennis_data_processed)

优化版本（适合大数据集，矢量化操作） #

如果你的数据集很大，iterrows和自定义apply函数效率较低，可以改用矢量化操作：

import pandas as pd

# 初始化示例数据（替换为你的真实DataFrame）
data = {
    'match_id': [1,2,3,4,5],
    'player1_id': ['A','B','A','A','B'],
    'player2_id': ['B','C','C','B','A'],
    'target': [1,1,0,0,1],
    'match_date': ['2023-01-01','2023-01-02','2023-01-03','2023-01-04','2023-01-05']
}
tennis_data_processed = pd.DataFrame(data)
tennis_data_processed['match_date'] = pd.to_datetime(tennis_data_processed['match_date'])

# 1. 拆分比赛记录为选手级记录
player_records = []
for _, row in tennis_data_processed.iterrows():
    player_records.append({
        'match_id': row['match_id'],
        'player_id': row['player1_id'],
        'win': 1 if row['target'] == 1 else 0,
        'match_date': row['match_date']
    })
    player_records.append({
        'match_id': row['match_id'],
        'player_id': row['player2_id'],
        'win': 0 if row['target'] == 1 else 1,
        'match_date': row['match_date']
    })

player_records_df = pd.DataFrame(player_records)
player_records_df = player_records_df.sort_values(by=['player_id', 'match_date']).reset_index(drop=True)

# 2. 矢量化计算streak
player_records_df['streak'] = 0
player_records_df['prev_win'] = player_records_df.groupby('player_id')['win'].shift(1)
player_records_df['prev_streak'] = player_records_df.groupby('player_id')['streak'].shift(1)

# 处理当前获胜的情况
win_mask = player_records_df['win'] == 1
player_records_df.loc[win_mask & (player_records_df['prev_win'] == 1), 'streak'] = player_records_df['prev_streak'] + 1
player_records_df.loc[win_mask & (player_records_df['prev_win'] == 0), 'streak'] = 1

# 处理当前失败的情况
lose_mask = player_records_df['win'] == 0
player_records_df.loc[lose_mask & (player_records_df['prev_win'] == 0), 'streak'] = player_records_df['prev_streak'] - 1
player_records_df.loc[lose_mask & (player_records_df['prev_win'] == 1), 'streak'] = -1

# 3. 合并回原DataFrame
streak_pivot = player_records_df.pivot(index='match_id', columns='player_id', values='streak')
tennis_data_processed['player1_streak'] = tennis_data_processed.apply(lambda row: streak_pivot.loc[row['match_id'], row['player1_id']], axis=1)
tennis_data_processed['player2_streak'] = tennis_data_processed.apply(lambda row: streak_pivot.loc[row['match_id'], row['player2_id']], axis=1)

# 查看结果
print(tennis_data_processed)

关键说明 #

1. 拆分记录：将每场比赛拆分为两个选手的独立记录，是为了单独追踪每个选手的每一场胜负，避免选手在player1或player2列切换导致的追踪困难。
2. 排序：必须按选手ID和比赛日期排序，确保streak的计算是按照比赛发生的先后顺序进行的，否则结果会出错。
3. streak计算逻辑：
   • 首次出场：streak为0
   • 当前获胜且上一场也获胜：streak = 上一场streak + 1
   • 当前获胜但上一场失败：streak重置为1
   • 当前失败且上一场也失败：streak = 上一场streak - 1
   • 当前失败但上一场获胜：streak重置为-1

注意事项 #

• 如果原DataFrame没有match_date列，需要确保数据本身是按比赛发生的先后顺序排列的，否则排序步骤需要调整。
• 选手ID必须是唯一标识（如选手编号），避免重名导致的计算错误。
• 大数据集优先使用矢量化版本，效率远高于循环遍历。

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