使用 pandas 根据条件选择滞后行和当前行答案

【问题标题】：Use pandas to select the lagged row along with current row based on criteria使用 pandas 根据条件选择滞后行和当前行
【发布时间】：2021-06-21 07:51:38
【问题描述】：

我有一个如下所示的数据框

person_id  source_system   r_diff
  1              A          NULL
  1              B           0
  1              B           9
  1              A           15
  1              A           574
  1              B           0
  1              A           63
  1              A           136
  1              B           0

我想根据 2 条规则的或运算来选择数据

a) 选择 source_system = B 的所有记录

b) 选择 n 和 n-1 行，其中 r_diff = 0。

例如，在上面的数据中，你可以发现 r_diff = 0 对应行号 2,6,9。所以，我想选择第 1,2 和 5,6 和 8,9 行。你可以看到我是如何选择 n 和 n-1 行的

我尝试了以下

df['flag_1'] = np.where((df['source_system'] == 'B'), '1','0')
df['flag_2'] = np.where((df['r_diff'] == 0), '1','0')
df['flag_3'] = np.where((df['r_diff'].shift(-1) == 0, '1','0')
df = df[((df['flag_1'] == '1') or (df['flag_2'] == '1') or (df['flag_3'] == '1'))]

我希望我的输出如下所示

person_id  source_system   r_diff
  1              A          NULL
  1              B           0
  1              B           9
  1              A           574
  1              B           0
  1              A           136
  1              B           0

【问题讨论】：

标签： python pandas dataframe numpy shift

【解决方案1】：

我认为你很接近，你可以将掩码设置为变量并通过| 链接按位OR 喜欢：

m1 = df['source_system'] == 'B'
m2 = df['r_diff'] == 0
m3 = df.groupby('person_id')['r_diff'].shift(-1) == 0

df = df[m1 | m2 | m3]
print (df)
   person_id source_system  r_diff
0          1             A     NaN
1          1             B     0.0
2          1             B     9.0
4          1             A   574.0
5          1             B     0.0
7          1             A   136.0
8          1             B     0.0

【讨论】：

我知道输出是正确的，但无论如何按person_id 分组并获得相同的输出？
@TheGreat - Editde 回答。