Python中按snp_id分组判断is_severe列生成YES/NO列问题

问题背景 #

原始数据（is_severe和snp_id均为字符串类型）：

is severe      sn_id
1                1
0                1
1                2
1                2

需求：按snp_id分组，若每组is_severe列同时包含至少一个"1"和"0"，新增列值为YES，否则为NO，预期结果：

is severe      sn_id     yes/no
    1                1      YES
    0                1      YES
    1                2      NO
    1                2      NO

错误代码分析 #

你使用的代码：

df['Value2']=np.where(df.groupby(df["snp_id"])['is_severe'].transform(lambda x: ((x=='1' ) & (x=="0")).any()), 'YES','NO')

核心问题在于(x=='1') & (x=="0")——这是判断同一行的is_severe同时等于"1"和"0"，逻辑上不可能成立，因此any()永远返回False，最终所有值都是NO。

修复后的代码 #

方法一：用集合判断组内是否包含目标值 #

import pandas as pd
import numpy as np

df['Value2'] = df.groupby('snp_id')['is_severe'].transform(
    lambda x: 'YES' if {'0', '1'}.issubset(x.unique()) else 'NO'
)

方法二：分别判断组内是否存在"1"和"0" #

df['Value2'] = df.groupby('snp_id')['is_severe'].transform(
    lambda x: 'YES' if (x == '1').any() and (x == '0').any() else 'NO'
)

方法三：先聚合再合并（适合大数据量） #

如果数据量较大，先分组聚合判断条件再合并回原数据，效率更高：

# 分组计算每组是否包含1和0
group_result = df.groupby('snp_id')['is_severe'].agg(
    has_1=lambda x: (x == '1').any(),
    has_0=lambda x: (x == '0').any()
).assign(Value2=lambda x: np.where(x['has_1'] & x['has_0'], 'YES', 'NO'))

# 合并回原数据
df = df.merge(group_result[['Value2']], on='snp_id', how='left')

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