在Pandas DataFrame中计算特定事件前事件数及特定事件统计

任务说明 #

• 任务1：针对每个ID，计算每个PUMP RUN事件发生前的其他事件累计数量并存入新列；ID切换时累计计数需重置。例如某ID下首次PUMP RUN前有5个其他事件，第二次前有3个。
• 任务2：统计每个ID对应的PUMP RUN事件总数。

现有代码问题 #

1. 任务1的代码未按ID正确分组，导致跨ID的计数混乱：

data['Count'] = data.groupby(data['EventDescription'].str.startswith('Pump RUN').cumsum()).cumcount().add(1)
data['EventsCount'] = data.groupby('ID')['Count'].apply(lambda x: x.shift(1))

2. 任务2已实现，但使用了中间列PumpRun，希望简化：

data['PumpRun'] = np.where(data['EventDescription'].str.startswith('Pump RUN'), 1, 0)
data['RunCount'] = data.groupby(['ID']).PumpRun.transform('sum')

修正后的完整代码 #

import pandas as pd

# 读取数据
data = pd.read_csv('eventdata.csv')
# 按指定字段排序（确保事件时序正确）
data = data.sort_values(by=["SiteNumber", "SerialNumber", "ID", "Time"], ascending=True)
# 过滤空ID行
data = data[data['ID'] != ''].reset_index(drop=True)

# 任务1：计算每个PUMP RUN前的其他事件累计数
# 标记PUMP RUN事件
data['is_pump_run'] = data['EventDescription'].str.startswith('Pump RUN')
# 按ID+PUMP RUN分组，在组内统计前置事件数
data['EventsCount'] = data.groupby(['ID', data['is_pump_run'].cumsum()]).apply(
    lambda x: x['is_pump_run'].cumsum().shift().fillna(0) if x['is_pump_run'].any() else pd.Series([0]*len(x))
).reset_index(drop=True)
# 非RUN事件行填充0
data['EventsCount'] = data['EventsCount'].fillna(0).astype(int)

# 任务2：简化统计每个ID的PUMP RUN总数（无需中间列）
data['RunCount'] = data.groupby('ID')['EventDescription'].transform(
    lambda x: x.str.startswith('Pump RUN').sum()
)

# 移除临时列
data.drop('is_pump_run', axis=1, inplace=True)
# 保存结果
data.to_csv("events_count.csv", index=False)

代码说明 #

• 任务1：通过ID+PUMP RUN累计分组的双重分组逻辑，确保每个ID内的计数独立，且每次PUMP RUN后自动重置计数；仅在PUMP RUN事件行显示前置累计数，非RUN事件行填充0。
• 任务2：直接在transform中对分组后的事件描述做条件求和，省去中间标记列，代码更简洁高效。

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