如何在Python DataFrame中筛选含多个指定前缀ID的集群行？

现有一个名为data的DataFrame，包含ID和cluster列，数据示例如下：

ID      cluster 
FP_101   1  
FP_102   1     
SP_209   3
SP_300   3
SP_209   1
FP_45    90
SP_50    90
FP_398   100
...

需求是筛选并输出存在至少一个前缀（FP或SP）对应多个不同ID的集群的所有行。现有两段代码，但不知道如何正确组合：

• 筛选以FP/SP开头的ID：data = data[data['ID'].str.startswith('FP')]（SP同理）
• 筛选集群的函数：data = data.groupby(['cluster']).filter(lambda x: x['ID'].nunique() > 1)

期望得到的结果如下：

ID      cluster 
    FP_101   1  
    FP_102   1
    SP_209   1     
    SP_209   3
    SP_300   3

从你的期望结果来看，正确的逻辑是先保留所有FP/SP开头的行，再筛选出那些FP前缀有多个不同ID，或者SP前缀有多个不同ID的集群，具体代码如下：

第一步：过滤出FP/SP开头的ID行 #

先把非FP/SP开头的ID行排除（如果数据集里有的话）：

filtered_ids = data[data['ID'].str.startswith(('FP', 'SP'))]

第二步：按集群筛选符合条件的分组 #

定义一个筛选函数，检查每个集群里FP或SP前缀的不同ID数量是否大于1，再用filter函数保留这些集群的所有行：

def check_cluster(group):
    # 统计集群内FP开头的不同ID数量
    fp_unique_count = group[group['ID'].str.startswith('FP')]['ID'].nunique()
    # 统计集群内SP开头的不同ID数量
    sp_unique_count = group[group['ID'].str.startswith('SP')]['ID'].nunique()
    # 满足任一条件则保留该集群
    return fp_unique_count > 1 or sp_unique_count > 1

# 应用筛选得到最终结果
result = filtered_ids.groupby('cluster').filter(check_cluster)

结果说明 #

运行后得到的结果和你的期望完全一致：

• cluster 1里有FP_101、FP_102两个不同的FP前缀ID，符合条件，保留所有行；
• cluster 3里有SP_209、SP_300两个不同的SP前缀ID，符合条件，保留所有行；
• cluster 90里FP和SP前缀各只有一个不同ID，不符合条件，被排除；
• cluster 100只有一个FP前缀ID，不符合条件，被排除。

为什么原有代码组合不行？ #

如果直接按你原来的代码顺序执行，会先只保留FP开头的行，漏掉SP开头的数据；如果分开筛选FP和SP再合并，逻辑会混乱且重复。上面的代码先统一保留所有目标前缀的行，再针对每个集群检查前缀对应的ID多样性，完美匹配需求。

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