groupby transform with if condition in pandas-CodePudding

I have a data frame as given below

df = pd.DataFrame({'key': ['a', 'a', 'a', 'b', 'c', 'c'] , 'val' : [10, np.nan, 9 , 10, 11, 13]})
df

key val
0   a   10.0
1   a   NaN
2   a   9.0
3   b   10.0
4   c   11.0
5   c   13.0

I want to perform groupby and transform that new column is each value divided by group mean , which I can do as below

df['new'] = df.groupby('key')['val'].transform(lambda g :  g/g.mean())
df.new

0    1.052632
1         NaN
2    0.947368
3    1.000000
4    0.916667
5    1.083333
Name: new, dtype: float64

Now I have condition that if val is np.nan then new column value will be np.inf which should result as below

0    1.052632
1    np.inf
2    0.947368
3    1.000000
4    0.916667
5    1.083333
Name: new, dtype: float64

In other words how can I have this check if a val is np.nan with groupby and transform.

Thanks in advance

CodePudding user response：

Add Series.replace:

df['new'] = (df.groupby('key')['val'].transform(lambda g :  g/g.mean())
               .replace(np.nan, np.inf))
print (df)
  key   val       new
0   a  10.0  1.052632
1   a   NaN       inf
2   a   9.0  0.947368
3   b  10.0  1.000000
4   c  11.0  0.916667
5   c  13.0  1.083333

Or numpy.where:

df['new'] = np.where(df.val.isna(), 
                     np.inf, df.groupby('key')['val'].transform(lambda g :  g/g.mean()))
print (df)
  key   val       new
0   a  10.0  1.052632
1   a   NaN       inf
2   a   9.0  0.947368
3   b  10.0  1.000000
4   c  11.0  0.916667
5   c  13.0  1.083333