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?
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
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.