如何用Pandas标记缺货y天后的x个符合条件的后续工作日？

在零售业务场景中，产品缺货恢复后的几日订单会异常偏高，为排除这些订单的影响，需要标记这些峰值日期。

初始DataFrame结构 #

已生成缺货累计列的代码 #

已通过以下代码计算出缺货天数累计列shortage_count：

rolling_func = lambda x: x.shift().rolling(2, min_periods=2).sum()
df['shortage_count'] = (df.assign(stock = df.stock.where(df.work_day == 1)).groupby('ITEM').stock.transform(rolling_func))

包含shortage_count的DataFrame #

需求说明 #

需要为每个产品标记出缺货时长≥y天后的x个满足work_day==1的后续行：

• 示例：当x=3、y=2时，需标记产品A的14/08/23、15/08/23、16/08/23
• 要求：创建flag列，对shortage_count≥2的行，为后续3个work_day==1的行设置flag=1，必须使用groupby和transform方法实现

解决方案 #

通过自定义transform函数实现逻辑，核心是在每个产品分组内识别触发点，然后为后续指定数量的工作日标记flag。

代码实现 #

def mark_post_shortage_days(group, x=3, y=2):
    # 初始化flag列表
    flag = [0] * len(group)
    # 找到缺货时长达标触发点的索引
    trigger_indices = group[group['shortage_count'] >= y].index
    if not trigger_indices.empty:
        # 取第一个触发点（若存在多个触发点，可根据业务需求调整逻辑）
        start_idx = trigger_indices[0]
        # 筛选触发点之后的工作日，取前x个
        target_days = group.loc[start_idx:][group['work_day'] == 1].head(x)
        # 为目标行设置flag=1
        for idx in target_days.index:
            flag[group.index.get_loc(idx)] = 1
    return flag

# 应用到每个产品分组，生成flag列
df['flag'] = df.groupby('Item').transform(mark_post_shortage_days, x=3, y=2)

代码说明 #

1. 自定义函数mark_post_shortage_days接收单个产品的分组数据，以及参数x（需标记的工作日数量）、y（缺货时长阈值）
2. 先定位分组内shortage_count≥y的触发行，再从该行开始筛选后续的工作日，取前x个进行标记
3. 通过groupby('Item').transform将函数批量应用到所有产品分组，自动生成flag列

执行结果 #

最终DataFrame新增flag列，符合示例要求的结果如下：

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