快速截取Pandas DataFrame：按ID保留至首次满足状态条件的行

针对你的需求——按id分组，保留每个分组从起始行到**首次出现status=1**的所有行（首次1的行保留，之后所有行不管status是什么都丢弃），这里提供一个比lambda x: x.cumsum().cumsum().le(1)快得多的矢量化方案：

步骤说明 #

1. 给每个分组内的行分配连续序号（从0开始）
2. 找出每个id分组中首次出现status=1的行序号
3. 筛选出每个分组内序号≤首次1行序号的所有行

代码实现 #

import pandas as pd

# 你的原始数据
d = {'id': [1,1,1,1,1,1,1,2,2,2,2,2,2,2], 'status': [0,0,0,0,1,1,1,0,0,0,0,1,0,1]}
df = pd.DataFrame(data=d)

# 1. 生成每个分组内的行号
df['group_row'] = df.groupby('id').cumcount()

# 2. 获取每个id首次出现status=1的分组内行号
first_1_row = df[df['status'] == 1].groupby('id')['group_row'].min()

# 3. 合并行号信息并筛选
df_merged = df.merge(first_1_row, on='id', suffixes=('', '_first'))
result = df_merged[df_merged['group_row'] <= df_merged['group_row_first']].drop(['group_row', 'group_row_first'], axis=1)

print(result)

输出结果 #

id  status
0   1       0
1   1       0
2   1       0
3   1       0
4   1       1
7   2       0
8   2       0
9   2       0
10  2       0
11  2       1

为什么更快？ #

这个方案全程用Pandas内置的矢量化操作（groupby.cumcount、groupby.min、merge），这些操作都是基于C语言底层实现的，比自定义lambda的apply快几个数量级，完全适配7万行的数据集。

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