DataFrame панд находится в следующем ряду

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