如何基于同组前置BRED事件日期为DataFrame新增CDAT列

需求说明 #

需要为DataFrame新增一列CDAT，具体规则：

• 仅当EVENT字段为PREG时填充CDAT，取值为同一ID、LACT、FDAT分组内，早于当前PREG事件DATE的最近一次BRED事件的DATE
• 非PREG事件的CDAT列保持为空

示例输入数据 #

ID  LACT    FDAT    EVENT   DATE
0   46  1   2011-09-23  BRED    2012-03-02
1   46  1   2011-09-23  PREG    2012-04-03
2   46  1   2011-09-23  PREG    2012-05-22
3   46  1   2011-09-23  PREG    2012-10-09
4   46  2   2012-11-15  FRESH   2012-11-15
5   46  2   2012-11-15  LUT     2013-01-08
6   46  2   2012-11-15  OS      2013-01-15
7   46  2   2012-11-15  BRED    2013-01-01
8   46  2   2012-11-15  BRED    2013-01-24
9   46  2   2012-11-15  PREG    2013-02-26
10  46  2   2012-11-16  BRED    2013-03-10

期望输出结果 #

ID  LACT    FDAT    EVENT   DATE         CDAT
0   46  1   2011-09-23  BRED    2012-03-02
1   46  1   2011-09-23  PREG    2012-04-03   2012-03-02
2   46  1   2011-09-23  PREG    2012-05-22   2012-03-02
3   46  1   2011-09-23  PREG    2012-10-09   2012-03-02
4   46  2   2012-11-15  FRESH   2012-11-15
5   46  2   2012-11-15  LUT     2013-01-08
6   46  2   2012-11-15  OS      2013-01-15
7   46  2   2012-11-15  BRED    2013-01-01
8   46  2   2012-11-15  BRED    2013-01-24
9   46  2   2012-11-15  PREG    2013-02-26   2013-01-24
10  46  2   2012-11-16  BRED    2013-03-10

完整样本数据列表 #

[[46,1,Timestamp('2011-09-23 00:00:00'),'BRED',Timestamp('2012-03-02 00:00:00')],
 [46,1,Timestamp('2011-09-23 00:00:00'),'PREG',Timestamp('2012-04-03 00:00:00')],
 [46,1,Timestamp('2011-09-23 00:00:00'),'PREG',Timestamp('2012-05-22 00:00:00')],
 [46,1,Timestamp('2011-09-23 00:00:00'),'PREG',Timestamp('2012-10-09 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'FRESH',Timestamp('2012-11-15 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'LUT',Timestamp('2013-01-08 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'OS',Timestamp('2013-01-15 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'BRED',Timestamp('2013-01-01 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'BRED',Timestamp('2013-01-24 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'PREG',Timestamp('2013-02-26 00:00:00')],
 [46,2,Timestamp('2012-11-16 00:00:00'),'BRED',Timestamp('2013-03-10 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'PREG',Timestamp('2013-04-16 00:00:00')],
 [46,2,Timestamp('2001-11-15 00:00:00'),'PREG',Timestamp('2013-08-06 00:00:00')]]

解决方案 #

可以通过groupby结合merge_asof实现，步骤如下：

1. 分离BRED和PREG事件的子数据集
2. 使用merge_asof按分组键匹配，找到每个PREG事件之前最近的BRED事件日期
3. 将匹配结果合并回原数据集，填充CDAT列

代码实现：

import pandas as pd
from pandas import Timestamp

# 构建样本DataFrame
data = [[46,1,Timestamp('2011-09-23 00:00:00'),'BRED',Timestamp('2012-03-02 00:00:00')],
 [46,1,Timestamp('2011-09-23 00:00:00'),'PREG',Timestamp('2012-04-03 00:00:00')],
 [46,1,Timestamp('2011-09-23 00:00:00'),'PREG',Timestamp('2012-05-22 00:00:00')],
 [46,1,Timestamp('2011-09-23 00:00:00'),'PREG',Timestamp('2012-10-09 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'FRESH',Timestamp('2012-11-15 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'LUT',Timestamp('2013-01-08 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'OS',Timestamp('2013-01-15 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'BRED',Timestamp('2013-01-01 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'BRED',Timestamp('2013-01-24 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'PREG',Timestamp('2013-02-26 00:00:00')],
 [46,2,Timestamp('2012-11-16 00:00:00'),'BRED',Timestamp('2013-03-10 00:00:00')],
 [46,2,Timestamp('2012-11-15 00:00:00'),'PREG',Timestamp('2013-04-16 00:00:00')],
 [46,2,Timestamp('2001-11-15 00:00:00'),'PREG',Timestamp('2013-08-06 00:00:00')]]

df = pd.DataFrame(data, columns=['ID', 'LACT', 'FDAT', 'EVENT', 'DATE'])

# 1. 分离BRED和PREG事件数据
bred_df = df[df['EVENT'] == 'BRED'].rename(columns={'DATE': 'BRED_DATE'})
preg_df = df[df['EVENT'] == 'PREG']

# 2. 对两个数据集按分组键和日期排序（merge_asof要求）
bred_df_sorted = bred_df.sort_values(by=['ID', 'LACT', 'FDAT', 'BRED_DATE'])
preg_df_sorted = preg_df.sort_values(by=['ID', 'LACT', 'FDAT', 'DATE'])

# 执行asof合并，匹配分组内日期<=当前PREG日期的最近BRED记录
matched = pd.merge_asof(
    preg_df_sorted,
    bred_df_sorted[['ID', 'LACT', 'FDAT', 'BRED_DATE']],
    on='DATE',
    by=['ID', 'LACT', 'FDAT'],
    direction='backward'
)

# 3. 合并回原数据集并整理CDAT列
df = df.merge(
    matched[['ID', 'LACT', 'FDAT', 'DATE', 'BRED_DATE']],
    on=['ID', 'LACT', 'FDAT', 'DATE'],
    how='left'
).rename(columns={'BRED_DATE': 'CDAT'})

# 清空非PREG事件的CDAT值
df.loc[df['EVENT'] != 'PREG', 'CDAT'] = None

print(df)

代码说明 #

• merge_asof是实现“最近日期匹配”的核心，direction='backward'确保只选取早于等于当前PREG日期的BRED记录，且取最近的一条
• 提前排序是merge_asof的强制要求，否则无法正确匹配
• 最后通过merge将匹配结果整合回原表，并对非PREG事件的CDAT值置空

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