如何在Pandas的groupby transform中添加np.nan判断条件？

原始数据 #

import pandas as pd
import numpy as np

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

对应的DataFrame输出：

原有实现 #

通过groupby.transform计算每个val除以对应key组的均值：

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

得到的new列结果：

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

需求 #

当val为np.nan时，将new列对应值设为np.inf，期望结果如下：

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

两种可行方案 #

方案一：在transform的lambda中直接判断 #

利用np.where在组内判断每个值是否为NaN，是则返回np.inf，否则计算除以组均值：

df['new'] = df.groupby('key')['val'].transform(
    lambda g: np.where(g.isna(), np.inf, g / g.mean())
)

方案二：先计算结果再替换NaN #

先按原有逻辑计算new列，再用fillna将其中的NaN替换为np.inf，写法更简洁直观：

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

内容的提问来源于stack exchange，提问作者Bharat Sharma