如何基于pandas DataFrame列值提取带上下扩展边界的子集

方案逻辑 #

完全基于Pandas内置向量化函数实现，无迭代遍历逻辑，自动处理边界避免索引越界，支持自定义上下扩展行数。

实现代码 #

import pandas as pd

# -------------------------- 可自定义参数 --------------------------
up_extend = 2    # 向上扩展行数
down_extend = 2  # 向下扩展行数
# -------------------------------------------------------------------

# 1. 生成SIGNAL=1的基础掩码
signal_mask = df['SIGNAL'].eq(1)

# 2. 生成向上扩展范围掩码：覆盖信号位前up_extend行到信号位
up_mask = signal_mask.rolling(up_extend + 1, min_periods=1).max().astype(bool)

# 3. 生成向下扩展范围掩码：覆盖信号位到信号位后down_extend行
down_mask = signal_mask[::-1].rolling(down_extend + 1, min_periods=1).max()[::-1].astype(bool)

# 4. 合并掩码提取结果，自动去重重叠覆盖的行
result_df = df[up_mask | down_mask].copy()

方案说明 #

• 边界处理：通过rolling的min_periods=1参数自动适配DataFrame首尾边界，不会出现索引越界报错
• 灵活性：仅需修改up_extend和down_extend的数值即可调整上下扩展的行数
• 性能：向量化操作远快于iterrows遍历，尤其适合十万行及以上的大数据量场景
• 自动去重：多个信号位的扩展范围重叠时，会自动保留唯一行，无需额外去重操作

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