在Pandas DataFrame中筛选ID重复且Location列存在不同值的所有对应行的最优方法

问题场景 #

你有如下的DataFrame：

需求是：保留所有ID对应的Location存在不同取值的行——只要某个ID下有至少两个不一样的Location值，就保留该ID的全部行（不管该ID出现2次还是3次以上），最终得到目标结果。

最优解决方案 #

用pandas的groupby结合transform方法是最高效的实现方式，代码如下：

import pandas as pd

# 构造你的原始DataFrame（如果已有数据可跳过这一步）
df = pd.DataFrame({
    'ID': [1, 1, 2, 3, 3, 4, 4, 4],
    'Location': ['Avenue', 'Block', 'Avenue', 'Street', 'Street', 'Av.', 'Street', 'Av.'],
    'Valid': ['Yes', 'No', 'Yes', 'No', 'No', 'Yes', 'No', 'Yes']
})

# 生成筛选掩码：每个ID对应的Location唯一值数量大于1的行标记为True
mask = df.groupby('ID')['Location'].transform('nunique') > 1
# 应用掩码得到结果
result = df[mask]

# 输出结果
print(result)

代码解释 #

1. groupby('ID')['Location']：按ID分组，聚焦Location列
2. transform('nunique')：对每个分组计算Location的唯一值数量，然后把这个结果广播到该分组的每一行（这样每一行都能拿到自己所属ID的Location唯一值数量）
3. > 1：生成布尔掩码，标记出那些所属ID的Location有不同值的行
4. df[mask]：用掩码筛选出符合条件的所有行

这个方法的优势：

• 高效：完全是pandas的向量化操作，比循环、自定义apply函数快得多，适合处理大型数据集
• 简洁：逻辑直观，一行代码生成掩码，一步完成筛选
• 通用：不管ID出现2次还是N次，只要存在不同的Location值，就会保留该ID的所有行

运行代码后得到的结果就是你期望的：

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