用groupby替代for循环计算玩家在试验中的总追踪时长

需要计算每个试验（trial）中各玩家（player）的总追踪时长。数据框包含['trial','time','player','event']列，其中event列记录track initiated、track update或track drop事件。现有for循环代码可完成计算，但面对50个试验、每个试验数百行数据时运行缓慢。希望移除for循环，改用groupby实现，同时处理并非所有玩家在试验结束时都会有track drop事件的情况。

原for循环代码：

import pandas as pd
df = pd.DataFrame({'trial':[101,101,101,101,101,101,101,101,101,102,102,102,102,102,102],
                   'time':[3,3,4,4,5,5,8,9,10,5,6,7,8,8,9],
                   'player':['tx-1','tx-2','tx-1','tx-2','tx-1','tx-2','tx-2','tx-2','tx-2','tx-1','tx-1','tx-1','tx-1','tx-3','tx-1'],
                   'event':['track initiated','track initiated','track update','track update','track drop','track drop',
'track initiated','track update','track drop','track initiated','track update','track update',
'track update','track initiated','track update']})

trials = set(df.trial)
track_results = []
for t in trials:
    trial_data = df[df['trial'] == t]
    unique_plrs = set(trial_data['player'])
    for plr in unique_plrs:
        plr_df = trial_data[trial_data['player'] == plr]
        plr_df.reset_index(inplace = True)
        total_track_time = 0
        for idx, row in plr_df.iterrows():
            time = row['time']
            event = row['event']
            if idx == 0:
                track_segment = 0
                prev_time = time
            elif 'initiated' in event:
                track_segment = 0
            elif 'update' in event:
                track_segment = time - prev_time
            elif 'drop' in event:
                track_segment = time - prev_time
            prev_time = time                     
            total_track_time += track_segment
        track_results.append([t,plr,total_track_time])  
track_results = pd.DataFrame(track_results,columns = ['trial','player','total track time'])           

核心思路是按trial和player分组，对每组内的时间序列计算有效追踪段的时长总和，跳过track initiated对应的时间段，同时处理未结束的追踪（无track drop的情况）。

步骤1：计算每组内的时间差 #

按trial和player分组，对每组的time列计算与前一行的差值（diff()方法），得到每行与上一行的时间间隔。

步骤2：标记有效追踪段 #

根据event列判断哪些时间间隔需要计入总时长：

• track initiated对应的间隔不计入（设为0）
• 每组第一行的间隔不计入（对应原循环中idx==0的情况，diff()返回NaN，需设为0）
• track update和track drop对应的间隔正常计入

步骤3：分组求和 #

对每组的有效时间间隔求和，得到每个玩家在对应试验中的总追踪时长。

完整优化代码：

import pandas as pd

# 示例数据
df = pd.DataFrame({
    'trial': [101,101,101,101,101,101,101,101,101,102,102,102,102,102,102],
    'time': [3,3,4,4,5,5,8,9,10,5,6,7,8,8,9],
    'player': ['tx-1','tx-2','tx-1','tx-2','tx-1','tx-2','tx-2','tx-2','tx-2','tx-1','tx-1','tx-1','tx-1','tx-3','tx-1'],
    'event': [
        'track initiated','track initiated','track update','track update','track drop','track drop',
        'track initiated','track update','track drop','track initiated','track update','track update',
        'track update','track initiated','track update'
    ]
})

# 1. 按trial和player分组，计算时间差
df['time_diff'] = df.groupby(['trial', 'player'])['time'].diff()

# 2. 处理需要排除的情况：第一行（diff为NaN）、track initiated事件
df['time_diff'] = df.apply(
    lambda row: 0 if pd.isna(row['time_diff']) or 'initiated' in row['event'] else row['time_diff'],
    axis=1
)

# 3. 分组求和得到总追踪时长
track_results = df.groupby(['trial', 'player'])['time_diff'].sum().reset_index()
track_results.rename(columns={'time_diff': 'total track time'}, inplace=True)

print(track_results)

结果验证 #

运行优化后的代码，输出结果与原for循环代码完全一致：

trial player  total track time
0    101   tx-1               2.0
1    101   tx-2               4.0
2    102   tx-1               4.0
3    102   tx-3               0.0

性能优势 #

相比嵌套for循环，groupby结合向量化操作避免了逐行迭代，数据量越大性能提升越明显，50个试验、数百行数据的场景下，运行速度可提升数倍甚至数十倍。

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