如何按ID分组标记Main_Date前后1个月内的Date？

需求说明 #

给定如下结构的DataFrame：

ID                Date  Main_Date
0   1 2015-04-02 12:00:00      False
1   1 2015-05-02 15:00:00       True
2   1 2015-02-15 00:00:00      False
3   2 2015-03-06 00:00:00       True
4   2 2016-11-20 00:00:00      False
5   3 2015-04-10 00:00:00       True
6   3 2016-07-28 00:00:00      False
7   3 2018-05-21 00:00:00      False

需要添加Within_Range布尔列，判断每个ID分组内的Date是否处于该组Main_Date=True对应的日期前后1个月范围内，预期输出如下：

ID                Date  Main_Date  Within_Range
0   1 2015-04-02 12:00:00      False          True
1   1 2015-05-02 15:00:00       True          True
2   1 2015-02-15 00:00:00      False         False
3   2 2015-03-06 00:00:00       True          True
4   2 2016-11-20 00:00:00      False         False
5   3 2015-04-10 00:00:00       True          True
6   3 2016-07-28 00:00:00      False         False
7   3 2018-05-21 00:00:00      False         False

当前代码存在的问题 #

用户编写的代码如下：

def check_if_within_range(date):
        # Select the `Main_Date` for the `ID` group 
        main_date = df.loc[df['Main_Date'] == True, 'Date'].iloc[0] #This line causes the problem
        # Specify the range of dates as all dates 1 month before and after the `Main_Date`            
        month_start = main_date - datetime.timedelta(days=30)
        month_end = main_date + datetime.timedelta(days=30)
        month = pd.date_range(month_start, month_end)
        # Check if `Date` is within range            
        return (date > month_start) & (date <= month_end)

df['Within_Range'] = df['Date'].apply(lambda x:check_if_within_range(x))
print(df)

存在两个核心问题：

• 未按ID分组获取对应主日期：代码直接提取整个DataFrame中第一个Main_Date=True的日期，无法匹配当前行所属的ID分组
• 日期范围处理不准确：用固定30天代替“1个月”不符合实际月天数差异，且pd.date_range的生成完全多余，逐行apply的效率极低

高效解决方案 #

以下是两种基于pandas向量操作的高效实现方式，均能解决上述问题：

方式一：使用groupby.transform广播主日期 #

import pandas as pd

# 确保Date列是datetime类型（如果原始数据是字符串的话）
df['Date'] = pd.to_datetime(df['Date'])

# 按ID分组，提取每个组的主日期（假设每个ID仅一个Main_Date=True的记录）
def get_group_main_date(group):
    return group.loc[group['Main_Date'], 'Date'].iloc[0]

# 给每个行添加对应组的主日期
df['_group_main_date'] = df.groupby('ID')['Date'].transform(get_group_main_date)

# 计算前后1个月的范围，用pd.DateOffset保证月天数的准确性
df['Within_Range'] = (
    df['Date'] >= df['_group_main_date'] - pd.DateOffset(months=1)
) & (
    df['Date'] <= df['_group_main_date'] + pd.DateOffset(months=1)
)

# 删除中间辅助列
df = df.drop('_group_main_date', axis=1)

print(df)

方式二：使用merge合并主日期后计算 #

import pandas as pd

df['Date'] = pd.to_datetime(df['Date'])

# 提取每个ID对应的主日期
main_date_map = df[df['Main_Date']][['ID', 'Date']].rename(columns={'Date': '_main_date'})

# 将主日期合并到原DataFrame的对应ID行
df = df.merge(main_date_map, on='ID', how='left')

# 使用between方法简化范围判断
df['Within_Range'] = df['Date'].between(
    df['_main_date'] - pd.DateOffset(months=1),
    df['_main_date'] + pd.DateOffset(months=1)
)

# 删除中间辅助列
df = df.drop('_main_date', axis=1)

print(df)

补充说明 #

• 两种方式均采用向量式计算，比逐行apply效率高10-100倍，适合处理大数据量
• pd.DateOffset(months=1)会自动处理不同月份的天数差异（比如2月、31天的大月），比固定30天更符合“1个月”的业务含义
• 如果需要将范围调整为覆盖整天（比如主日期是2015-05-02 15:00:00，希望结束到2015-06-02 23:59:59），可以将结束时间改为：

  df['_main_date'] + pd.DateOffset(months=1) - pd.Timedelta(seconds=1)
  

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