如何用Python/Pandas统计Protein对应不同Organism的数量并格式化输出

问题需求 #

现有多列CSV文件，其中Protein和Organism列存在大量重复值，需要完成以下操作：

1. 按Protein分组，统计每个Protein对应的不同Organism的数量
2. 将同组的ID收集为列表，Organism去重后收集为列表，date保留组内第一条记录
3. 生成指定格式的输出，并找出对应Organism种类最多的Protein

输入数据 #

ID   Protein   Organism            date
12   Prot      Homo sapiens        24-dec
ab1  domain    Mus musculus        12-Apr
14   Prot      Homo sapiens        15-Jun
ijk3 ring      Rattus              25-May
ghl  ring      Homo sapiens        23-Jul
cdk8 Prot      Gallus gallus       18-sep
23bg Prot      Eschereria coli     13-sep

期望输出 #

ID             Protein   Organism                                        date   count
[12,cdk8,23bg] Prot      [Homo sapiens, Gallus gallus, Eschereria coli]  24-dec  3
[ab1]          domain    [Mus musculus]                                  12-Apr  1
[ijk3, ghl]    ring      [Rattus, Homo sapiens]                          25-May  2

现有代码问题分析 #

你尝试的代码存在两个核心问题：

1. CSV读取分隔符错误：输入数据是用多个空格分隔的，但代码中指定了delimiter=','，导致数据读入时列识别错误，后续分组聚合自然无法得到正确结果
2. 聚合逻辑不匹配需求：直接用agg(set)会将所有列转为集合，既无法保留date的第一条记录，也没有单独生成count统计列，且集合是无序的，与期望输出的列表格式不符

修正后的解决方案代码 #

import pandas as pd

# 正确读取空格分隔的CSV文件
df = pd.read_csv('prot_info.csv', sep='\s+')

# 按Protein分组，指定各列的聚合规则
df_out = df.groupby('Protein').agg(
    ID=('ID', list),
    Organism=('Organism', lambda x: list(x.unique())),
    date=('date', lambda x: x.iloc[0])  # 保留组内第一条日期
).reset_index()

# 添加count列，统计每个Protein对应的不同Organism数量
df_out['count'] = df_out['Organism'].apply(len)

# 调整列顺序，匹配期望输出
df_out = df_out[['ID', 'Protein', 'Organism', 'date', 'count']]

# 找出Organism种类最多的Protein
max_organism_prot = df_out.loc[df_out['count'].idxmax()]
print("对应Organism种类最多的Protein：")
print(max_organism_prot)

# 保存结果到TSV文件
df_out.to_csv('test.csv', sep='\t', index=False)

代码说明 #

1. 读取CSV：使用sep='\s+'匹配任意数量的空格，正确识别各列
2. 分组聚合：
   • ID列：将同组的所有ID收集为列表
   • Organism列：用x.unique()去重后转为列表，保留出现顺序
   • date列：用x.iloc[0]取组内第一条记录的日期
3. 统计数量：通过apply(len)获取每个Organism列表的长度，即不同种类的数量
4. 找出最多种类的Protein：利用idxmax()找到count列最大值对应的行

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