熊猫DataFrame isin和下一行

Question

对于给定的DataFrame，按b排序并重置索引：

df = pd.DataFrame({'a': list('abcdef'), 
                   'b': [0, 2, 7, 3, 9, 15]} 
                 ).sort_values('b').reset_index(drop=True)
   a   b
0  a   0
1  b   2
2  d   3
3  c   7
4  e   9
5  f  15

和一个列表，v

v = list('adf')

我只想提取v中的行和接下来的行（如果有的话），类似于grep -A1：

   a   b
0  a   0
1  b   2
2  d   3
3  c   7
5  f  15

我可以通过将isin的索引和isin的索引加一个连接起来来做到这一点，就像这样：

df[df.index.isin(
    np.concatenate(
        (df[df['a'].isin(v)].index, 
         df[df['a'].isin(v)].index + 1)))]

但这很长，而且不太容易理解。有更好的方法吗？

Answer 1

您可以结合使用isin条件和shift（下一行）来创建所需的布尔值：

df[df.a.isin(v).pipe(lambda x: x | x.shift())]

#   a   b
#0  a   0
#1  b   2
#2  d   3
#3  c   7
#5  f  15