Keep rows and neighbour rows that satisfy condition

user_6512118 0
Viewed 56

For a DataFrame such as

import pandas as pd
import numpy as np
times = [21 , 34, 37, 40, 55, 65, 67, 84, 88, 90 , 91, 97, 104,105, 108]
names = ['bob', 'alice', 'bob', np.NaN , 'ali', 'alice', np.NaN , 'ali', 'moji', 'ali', 'moji', np.NaN , 'bob', 'bob', 'bob']
actions = ['enter', 'enter', 'search', 'search', 'enter', 'search', 'purchase', 'exit', 'enter' , 'enter', 'search', 'purchase', 'exit', 'enter', 'purchase']
df = pd.DataFrame({'name' : names , 'action': actions, 'time' : times})

I want only rows where name is NaN and also row exactly before and exactly after such rows. I could do this by for and if commands. but is there any more neat way to do this?

pandas numpy python
2 Answers

Use Series.isna then Series.shift to get the the row before and after:

s1 = df['name'].isna()
s2 = s1.shift()
s3 = s1.shift(-1)

df[s1 | s2 | s3]

     name    action  time
2     bob    search    37
3     NaN    search    40
4     ali     enter    55
5   alice    search    65
6     NaN  purchase    67
7     ali      exit    84
10   moji    search    91
11    NaN  purchase    97
12    bob      exit   104
user_90812670

Obtain a mask where 'name' is NaN and convolve with a size 3 window:

ix_na = df['name'].isna().to_numpy()
m = np.convolve(ix_na, np.ones(3), mode='same').astype(bool)
print(df[m])

     name    action  time
2     bob    search    37
3     NaN    search    40
4     ali     enter    55
5   alice    search    65
6     NaN  purchase    67
7     ali      exit    84
10   moji    search    91
11    NaN  purchase    97
12    bob      exit   104

Or we can also use Series.rolling:

df[df['name'].isna().rolling(3, min_periods=0, center=True).sum().gt(0)]
user_96986840
Related