Q:

Keep NaN values using apply and lambda in python

I am working with python and I have a list of value in column A including some NaN values.
I need to code them as 1 and 0 depending if they are above or below the median (0.4). I've tried the code below which is working fairly well
df["A_median"] = (df["A"].apply(lambda count: 0 if count<median_a else 1))

But it is giving me the NaN values as 1 instead of keeping them. 
  A   A_median
36.6  1
NaN   1
NaN   1
0.1   0

The final table instead should be as follow
  A   A_median
36.6  1
NaN   NaN
NaN   NaN
7     0

Any idea on how I can fix this?
Thanks

A:

Is it:
median_a = 0.4
df['A_median'] = (df['A'] > median_a).where(df.A.notna())

Output:
      A  A_median
0  36.6       1.0
1   NaN       NaN
2   NaN       NaN
3   0.1       0.0

