如何在Pandas中当分组内列nunique大于N时修改对应列值？

原始数据 #

import pandas as pd
df = pd.DataFrame(
    {
        'a': ['a', 'a', 'a', 'b', 'c', 'x', 'j', 'w'],
        'b': [1, 1, 1, 2, 2, 3, 3, 3],
    }
)

预期输出 #

修改列a后的数据：

a  b  
0    a  1  
1    a  1  
2    a  1  
3  NaN  2  
4  NaN  2  
5  NaN  3  
6  NaN  3  
7  NaN  3  

逻辑规则 #

以列b进行分组，若分组内a列的唯一值数量大于1，则将该分组内a列的所有值设为np.nan。

现有实现 #

已有的可正常运行的代码：

import numpy as np
df['x'] = df.groupby('b')['a'].transform('nunique')
df.loc[df.x > 1, 'a'] = np.nan

更简洁高效的实现 #

可以省去临时列的创建，直接通过分组变换完成修改，以下两种写法更紧凑且内存效率更高：

写法一：lambda表达式结合transform #

import numpy as np
df['a'] = df.groupby('b')['a'].transform(lambda g: g if g.nunique() == 1 else np.nan)

写法二：where方法结合分组变换 #

import numpy as np
df['a'] = df['a'].where(df.groupby('b')['a'].transform('nunique') == 1, np.nan)

这两种方法无需额外生成临时列，直接对目标列a做修改，在数据量较大时能减少内存占用，代码逻辑也更直观。

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