python 3.x - Get next non nan values in pandas dataframe

Question

I have a data frame like below:

A    B  
10  NaN
NaN 20
NaN NaN
NaN NaN
NaN NaN
NaN 50
20  NaN
NaN 30
NaN NaN
30  30
40  NaN
NaN 10

Here I need to return previous and next B column value for each non NaN values of column A.

The code which I'm using is:

df['prev_b'] = NP.where(df['A'].notna(), df['B'].shift(-1),NP.nan)

df['next_b'] = NP.where(df['A'].notna(), df['B'].shift(1),NP.nan)

The required output is:

A    B   prev_b next_b
10  NaN  NaN    20
NaN 20   NaN    NaN
NaN NaN  NaN    NaN
NaN NaN  NaN    NaN
NaN NaN  NaN    NaN
NaN 50   NaN    Nan
20  NaN  50     30  
NaN 30   NaN    NaN
NaN NaN  NaN    NaN
30  30   30     30
40  NaN  30     10
NaN 10   NaN    NaN

Someone help me in correcting my logic.

question from:https://stackoverflow.com/questions/65924897/get-next-non-nan-values-in-pandas-dataframe

Answer 1

Use a forward or backward fill instead in your numpy where; it should correctly align to get your next/previous non-nan value:

    df.assign(
    prev_b=np.where(df.A.notna(), df.B.ffill(), np.nan),
    next_b=np.where(df.A.notna(), df.B.bfill(), np.nan),
)