pandas DataFrame isin and following row

user_3279716 0
Viewed 417

For a given DataFrame, sorted by b and index reset:

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

and a list, v

v = list('adf')

I would like to pull out just the rows in v and the following row (if there is one), similar to grep -A1:

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

I can do this by concatenating the index from isin and the index from isin plus one, like so:

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

But this is long and not too easy to understand. Is there a better way?

pandas python
1 Answers

You can combine the isin condition and the shift (next row) to create the boolean you needed:

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
user_49834500
Related