Pandas技术实现：按Session分组移除无前置START的首个END值

原始数据构造 #

先给出构造测试数据的代码：

import pandas as pd

data = {
    'Session': [1, 1, 2, 2, 3, 3, 4, 4, 5, 5],
    'Task': [pd.NA, 'END', 'START', 'END', 'END', pd.NA, pd.NA, pd.NA, pd.NA, 'START']
}
df = pd.DataFrame(data)

对应的初始数据表格：

预期处理结果 #

需要将分组内首个有效Task为END且该END前无START的行替换为NaN，预期结果如下：

实现方案 #

方案1：基于分组标记的高效处理（适合大数据量）

通过生成辅助列标记分组内非空Task的顺序，再筛选出需要替换的行：

# 生成辅助列：标记分组内非空Task的顺序（空值为-1，首个非空为0，依次递增）
df['non_empty_seq'] = df.groupby('Session')['Task'].apply(
    lambda x: x.notna().cumsum() - 1
)

# 筛选需要替换的行：Task为END且是分组内首个非空Task
replace_mask = (df['Task'] == 'END') & (df['non_empty_seq'] == 0)

# 执行替换
df.loc[replace_mask, 'Task'] = pd.NA

# 清理辅助列
df = df.drop('non_empty_seq', axis=1)

方案2：基于GroupBy Apply的直观处理（适合小数据量）

直接对每个分组单独处理，逻辑更清晰：

def clean_group(group):
    # 提取分组内非空的Task条目
    non_empty_tasks = group['Task'].dropna()
    if not non_empty_tasks.empty:
        # 首个有效Task是END时，替换对应行的Task为NaN
        if non_empty_tasks.iloc[0] == 'END':
            first_end_idx = non_empty_tasks.index[0]
            group.loc[first_end_idx, 'Task'] = pd.NA
    return group

# 应用分组处理
df = df.groupby('Session', group_keys=False).apply(clean_group)

结果验证 #

运行任意一种方案后，打印df即可得到符合预期的结果：

Session   Task
0        1   <NA>
1        1   <NA>
2        2  START
3        2    END
4        3   <NA>
5        3   <NA>
6        4   <NA>
7        4   <NA>
8        5   <NA>
9        5  START

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