如何在DataFrame中输出指定值所在行的上下N行（含目标行）

我的某列中存在以数值777表示的分类变量，我希望在识别到该列中的777时，输出其所在行的上下3行数据。目前我已实现输出上下1行及目标行的功能，代码如下，请问如何修改以实现输出3行或任意行数的上下行？

m = df['ResetPoint'] == 777
df = df[m.shift(-1)|m.shift()|m]

你可以通过定义变量n指定需要保留的上下行数，生成包含目标行及上下n行的掩码，最终筛选出这些行，核心逻辑是把目标行、向上偏移1到n行、向下偏移1到n行的所有行标记出来。

方法一：循环生成掩码 #

import pandas as pd

n = 3  # 可替换为任意需要的行数，比如1、5等
m = df['ResetPoint'] == 777
# 初始化全为False的掩码
mask = pd.Series(False, index=df.index)
# 遍历偏移量，标记所有符合条件的行
for k in range(-n, n + 1):
    mask |= m.shift(k, fill_value=False)
# 筛选结果
df_result = df[mask]

方法二：简洁的concat+any写法 #

import pandas as pd

n = 3
m = df['ResetPoint'] == 777
# 生成所有偏移后的掩码并取逻辑或
mask = pd.concat([m.shift(k, fill_value=False) for k in range(-n, n+1)], axis=1).any(axis=1)
df_result = df[mask]

说明 #

你原来的代码本质是n=1的情况，对应偏移量-1、0、1，也就是m.shift(-1)（向下1行）、m（目标行）、m.shift()（向上1行），和上面的逻辑完全一致。修改n的值就能灵活调整需要保留的上下行数。

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