python - 在 pandas 数据框中第一次出现条件后删除所有行

Question

使用以下示例:

df = pd.DataFrame({"Person":[1,1,2,2,2,2,3,3,3], "Bank":["OPEN","OPEN","OPEN","OPEN","CLOSED","OPEN","OPEN","CLOSED","CLOSED"]})

   Person   Bank
0       1   OPEN
1       1   OPEN
2       2   OPEN
3       2   OPEN
4       2   CLOSED
5       2   OPEN
6       3   OPEN
7       3   CLOSED
8       3   CLOSED

我想生成一个输出，使每个 Person 组的所有行都保持到并包括第一次出现 CLOSED 。所以它应该看起来像:

   Person   Bank
0       1   OPEN
1       1   OPEN
2       2   OPEN
3       2   OPEN
4       2   CLOSED
6       3   OPEN
7       3   CLOSED

我能够使用 this question构建一个接近的输出:

mask = (df['Bank']
    .where(df['Bank'] == 'OPEN')
    .groupby(df['Person'])
    .ffill(limit=1)
)
df[mask.notnull()]

# The above produces this
   Person   Bank
0       1   OPEN
1       1   OPEN
2       2   OPEN
3       2   OPEN
4       2   CLOSED
5       2   OPEN
6       3   OPEN
7       3   CLOSED

所以我当前的代码无法处理从CLOSED 回到OPEN 的情况。有没有一种不会非常慢的好方法？

Answer 1

您可以使用 groupby 创建掩码。它需要 2 个操作 cummax + shift 所以直接的方法是使用较慢的应用，但是对于许多组你会看到 2 个单独的 groupby 的性能更好 使用内置操作调用

m = (df['Bank'].eq('CLOSED')
       .groupby(df['Person'])
       .apply(lambda x: ~x.cummax().shift().fillna(False)))

# or
m = ~(df['Bank'].eq('CLOSED')
        .groupby(df['Person']).cummax()
        .groupby(df['Person']).shift()
        .fillna(False))

df[m]
   Person    Bank
0       1    OPEN
1       1    OPEN
2       2    OPEN
3       2    OPEN
4       2  CLOSED
6       3    OPEN
7       3  CLOSED