第一次满足条件时的子集 pandas 数据帧

Subset pandas dataframe up to when condition is met the first time

我没有运气完成任务,我想将 pandas 数据框子集化为一个值,并按其 id 分组。在实际数据集中,我在 'id' 和 'status'

之间有几列

例如:

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)

    id  status
0    1       0
1    1       0
2    1       0
3    1       0
4    1       1
5    1       1
6    1       1
7    2       0
8    2       0
9    2       0
10   2       0
11   2       1
12   2       0
13   2       1

所需的子集为:

    id  status
0    1       0
1    1       0
2    1       0
3    1       0
4    1       1
5    2       0
6    2       0
7    2       0
8    2       0
9    2       1

让我们试试 groupby + cumsum:

df = df.groupby('id', group_keys=False)\
       .apply(lambda x: x[x.status.cumsum().cumsum().le(1)])\
       .reset_index(drop=1)
df

   id  status
0   1       0
1   1       0
2   1       0
3   1       0
4   1       1
5   2       0
6   2       0
7   2       0
8   2       0
9   2       1

这是一个替代方法,它执行 groupby 来创建一个用作索引器的掩码:

df = df[df.status.eq(1).groupby(df.id)\
          .apply(lambda x: x.cumsum().cumsum().le(1))]\
          .reset_index(drop=1)
df

   id  status
0   1       0
1   1       0
2   1       0
3   1       0
4   1       1
5   2       0
6   2       0
7   2       0
8   2       0
9   2       1