Applying lambda to whole dataframe with if condition

I have a df that looks like this:

A
0
1
2
2.5
0
9.5
0

I want to calculate the mean of the columns where A>0 so that my df would look like this:

A       B
0       0
1       3.75
2       3.75
2.5     3.75
0       0
9.5     3.75
0       0

I use:

df.applymap(lambda x: x[x['A']>0].mean())

But get:

TypeError: ‘float’ object is not subscriptable

I also tried

df.apply(lambda x: x[x['A']>0].mean(), axis=1)

But get:

KeyError: False

Which is produced by the x['A']>0 mask.

And:

df.apply(lambda x: x[x['A']>0].mean() if x['A']>0 else 0, axis=1)

I couldn’t find a solution how can I apply a lambda function to a filtered dataframe. I could write a function, but wanted to know if its possible with apply(lambda).

Answer

Alternately you could use groupby and transform like

df['B'] = df.groupby(df['A'].gt(0)).transform('mean')
df

     A    B
0   0.0 0.00
1   1.0 3.75
2   2.0 3.75
3   2.5 3.75
4   0.0 0.00
5   9.5 3.75
6   0.0 0.00