Pandas标记DataFrame行：ID存在于指定年份集且不在另一集合

实现逻辑 #

打标规则是ID维度的判断，不需要逐行计算，核心判断条件为单个ID同时满足：

• 存在年份属于[2017,2018,2019]集合的记录
• 不存在任何年份属于[2020,2021,2022]集合的记录
  满足以上两个条件的ID，其下所有行label设为1，否则设为0。

实现代码 #

直接使用groupby + transform完成分组判断和结果广播，代码简洁且运行效率高：

import pandas as pd

# 示例数据
df1 = pd.DataFrame({
    'ID': ['AX1', 'AX1', 'AX1','AX1','AX1','AX1','AX2','AX2','AX2','AX3','AX3','AX4','AX4','AX4'],
    'year':[2017,2018,2019,2020,2021,2022,2019,2020,2022,2019,2020,2017,2018,2019]
})

# 定义规则对应的年份集合
early_years = {2017, 2018, 2019}
late_years = {2020, 2021, 2022}

# 分组计算标签并广播到所有行
df1['label'] = df1.groupby('ID')['year'].transform(
    lambda year_series: 1 if (set(year_series) & early_years) and not (set(year_series) & late_years) else 0
)

结果验证 #

运行代码后输出的DataFrame与期望结果完全一致：

ID  year  label
0   AX1  2017      0
1   AX1  2018      0
2   AX1  2019      0
3   AX1  2020      0
4   AX1  2021      0
5   AX1  2022      0
6   AX2  2019      0
7   AX2  2020      0
8   AX2  2022      0
9   AX3  2019      0
10  AX3  2020      0
11  AX4  2017      1
12  AX4  2018      1
13  AX4  2019      1

注：如果数据量极大，可以先对ID和year做去重后判断，再映射回原表，进一步提升运行速度。

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