如何在Pandas中基于日期与Flag条件删除DataFrame行？

现有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame([["A","13-02-2022","B","FALSE"],["A","13-02-2022","C","FALSE"],["A","14-02-2022","D","FALSE"],
                   ["A","14-02-2022","E","FALSE"],["A","16-02-2022","A","TRUE"],["A","16-02-2022","F","FALSE"],
                   ["A","17-02-2022","G","FALSE"],["A","17-02-2022","H","FALSE"],["A","18-02-2022","I","FALSE"],
                   ["A","18-02-2022","J","FALSE"]],columns=["id1","date","id2","flag"])

对应的表格：

需要实现的逻辑：

• 删除所有flag为TRUE的日期对应的所有行
• 同时删除该日期的前一个工作日的所有行
• 以及该日期的后一个工作日的所有行
• 特殊情况：如果TRUE出现在当月最后一天，且无后续工作日，则仅删除当天和前一个工作日的行

本例中，2月16日存在flag=TRUE，因此需要删除2月14日（前一个工作日）、2月16日、2月17日（后一个工作日）的所有行，预期输出如下：

df_out = pd.DataFrame([["A","13-02-2022","B","FALSE"],["A","13-02-2022","C","FALSE"],["A","18-02-2022","I","FALSE"],
                       ["A","18-02-2022","J","FALSE"]],columns=["id1","date","id2","flag"])

对应的表格：

可以按照以下步骤实现：

1. 转换日期格式并处理基础数据
   先把date列转换为datetime类型，方便后续工作日计算，同时提取排序后的唯一日期：

   # 转换日期列格式为datetime
   df['date'] = pd.to_datetime(df['date'], format='%d-%m-%Y')
   # 获取排序后的唯一日期集合
   unique_dates = df['date'].sort_values().unique()
   

2. 标记所有需要删除的日期
   先定位所有flag=TRUE的日期，再对每个目标日期计算前后工作日，将这些日期加入待删除集合：

   # 提取所有flag为TRUE的日期
   true_dates = df[df['flag'] == 'TRUE']['date'].unique()
   
   # 初始化待删除日期集合
   to_drop_dates = set()
   
   for dt in true_dates:
       # 添加当前日期
       to_drop_dates.add(dt)
       
       # 计算并添加前一个工作日
       prev_bday = pd.bdate_range(end=dt, periods=2)[0]
       to_drop_dates.add(prev_bday)
       
       # 判断是否为当月最后一天，再决定是否添加后一个工作日
       is_month_last_day = (dt + pd.Timedelta(days=1)).month != dt.month
       if not is_month_last_day:
           next_bday = pd.bdate_range(start=dt, periods=2)[1]
           # 确保后一个工作日仍在当月内
           if next_bday.month == dt.month:
               to_drop_dates.add(next_bday)
   

3. 过滤得到最终结果
   筛选出不在待删除日期列表中的行，可选恢复原日期字符串格式：

   # 过滤掉待删除日期的行
   df_out = df[~df['date'].isin(to_drop_dates)].copy()
   # 恢复日期为原字符串格式（可选操作）
   df_out['date'] = df_out['date'].dt.strftime('%d-%m-%Y')
   

执行上述代码后，df_out即为符合要求的结果。

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