基于reference列nunique过滤DataFrame，求特殊规则下的高效实现

问题背景 #

现有如下Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(list(zip(['name1', 'name2', 'name2', 'name3', 'name2', 'name4', 'name4'],
                        ['a', 'b', 'c', 'b', 'c', 'b', 'b'],
                       ['USA', 'USA', 'USA', 'USA', 'USA', 'USA', 'USA'],
                       [np.nan, np.nan, 'Jan', np.nan, np.nan, 'Jan', 'Jan'],
                       [1,2,3,4,5,6,7])),
                columns=['sample ID', 'reference', 'country', 'month', 'value'])

原本使用以下代码过滤：

df = df[df.groupby(['sample ID'])['reference'].transform('nunique') > 1]

现在需要新增规则：若某个sample ID对应的reference唯一值列表包含'a'，即便仅存在'a'这一个值，也判定其nunique>1，要求更高效的实现方式，避免拆分拼接的繁琐操作。

高效解决方案 #

可以通过一次分组计算同时满足两个条件：原有的nunique>1，或者分组内存在'a'。使用transform生成布尔掩码即可：

# 计算每个sample ID的过滤条件：nunique>1 或 包含'a'
mask = df.groupby('sample ID')['reference'].transform(
    lambda x: x.nunique() > 1 or (x == 'a').any()
)

# 应用掩码得到过滤后的DataFrame
filtered_df = df[mask]

逻辑解释 #

• x.nunique() >1：保留原规则中reference有多个唯一值的分组
• (x == 'a').any()：新增规则，只要分组内存在'a'就保留该分组的所有行
• 通过transform将分组级的判断结果广播到每一行，生成布尔掩码后直接过滤

这种方法仅需一次分组操作，避免了拆分拼接的额外开销，在大数据集上的执行效率更高。

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