python - 子集 pandas 数据框直到第一次满足条件时

Question

我没有运气完成一项任务，我想将 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

Answer 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