如何在Pandas中高效提取300码与前序200码间的行并添加批次列

原始数据 #

import pandas as pd

data = {
    'c1': [1, 1, 1, 2, 6, 6, 9, 1, 9, 8, 1, 2, 8, 9],
    'c2': [2, 5, 2, 4, 1, 1, 8, 6, 2, 1, 2, 5, 0, 1],
    'c3': [3, 7, 3, 1, 1, 5, 2, 9, 1, 2, 1, 3, 4, 2],
    'code': [200, 220, 200, 340, 370, 270, 300, 700, 200, 400, 200, 900, 300, 620]
}
df = pd.DataFrame(data)

需求说明 #

找到每一行code=300的记录，定位其**最近的前序code=200**的记录，提取这两个标记行之间的所有行（包含标记行本身），并添加batch_num列区分不同批次。

目标结果 #

c1  c2  c3  code  batch_num
0   1   2   3   200          0
1   2   4   1   340          0
2   6   1   1   370          0
3   6   1   5   270          0
4   9   8   2   300          0
5   1   2   1   200          1
6   2   5   3   900          1
7   8   0   4   300          1

实现步骤 #

1. 标记关键行并定位批次端点 #

先提取code=200和code=300的行索引，再为每个code=300行匹配最近的前序code=200行：

# 提取两类关键行的索引
idx_200 = df[df['code'] == 200].index
idx_300 = df[df['code'] == 300].index

# 用searchsorted快速匹配（大数据集更高效）
matched_200_idx = idx_200[idx_200.searchsorted(idx_300, side='right') - 1]

# 生成每个批次的索引区间：(起始200索引, 结束300索引)
batch_ranges = list(zip(matched_200_idx, idx_300))

2. 为行分配批次号 #

遍历每个批次区间，为区间内的行设置对应批次编号：

# 初始化批次列为空值
df['batch_num'] = pd.NA

for batch_num, (start, end) in enumerate(batch_ranges):
    df.loc[start:end, 'batch_num'] = batch_num

3. 筛选并整理结果 #

过滤出有效批次行，重置索引并调整数据类型：

# 筛选有批次号的行，重置索引
result = df[df['batch_num'].notna()].reset_index(drop=True)
# 转换批次号为整数类型
result['batch_num'] = result['batch_num'].astype(int)

print(result)

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