如何在Pandas DataFrame中提取每个BUY出现后的3行数据切片？

实现步骤 #

1. 定位DataFrame中Action列值为"BUY"的所有行索引
2. 遍历每个索引，从该位置开始提取最多3行的连续切片（自动处理末尾不足3行的情况）

代码示例 #

import pandas as pd

# 构造示例DataFrame
data = {
    'Time': ['1:00', '1:01', '1:02', '1:03'],
    'Price': [33, 35, 34, 35],
    'MA': [34, 36, 35, 33],
    'Action': ['BUY', 'BUY', 'BUY', 'SELL']
}
df = pd.DataFrame(data)

# 获取所有"BUY"对应的索引列表
buy_indices = df[df['Action'] == 'BUY'].index.tolist()

# 遍历每个索引提取切片
for idx in buy_indices:
    # loc为闭区间，取从idx到idx+2的行，超出范围自动取到DataFrame末尾
    slice_df = df.loc[idx:idx+2]
    print(f"从索引{idx}开始的切片：")
    print(slice_df)
    print("-" * 30)

输出结果 #

从索引0开始的切片：
   Time  Price  MA Action
0  1:00     33  34    BUY
1  1:01     35  36    BUY
2  1:02     34  35    BUY
------------------------------
从索引1开始的切片：
   Time  Price  MA Action
1  1:01     35  36    BUY
2  1:02     34  35    BUY
3  1:03     35  33   SELL
------------------------------
从索引2开始的切片：
   Time  Price  MA Action
2  1:02     34  35    BUY
3  1:03     35  33   SELL
------------------------------

可选优化：只保留长度为3的切片 #

如果需要过滤掉不足3行的切片，可添加判断逻辑：

for idx in buy_indices:
    # 确保索引+2不超过DataFrame的最大索引
    if idx + 2 <= df.index.max():
        slice_df = df.loc[idx:idx+2]
        print(f"从索引{idx}开始的切片：")
        print(slice_df)
        print("-" * 30)
    else:
        print(f"索引{idx}起始的切片不足3行，跳过")

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