如何在Pandas中对DataFrame指定列分组并合并同基因药物引用

需求说明 #

将DataFrame中同一基因对应同一药物的行合并，保留该药物的全部引用信息，非重复行保持不变。

示例输入 #

Gene                          Drug                     ID     Cite
1  MAD1L1                       Lithium[17]           34718328     [17]
2    OAS1                       Lithium[17]           34718328     [17]
3    OAS1                       Lithium[7]            27401222      [7]

期望输出 #

Gene                          Drug                     ID     Cite
1  MAD1L1                       Lithium[17]           34718328     [17]
2    OAS1                       Lithium[17][7]        34718328     [17]

实现代码（Pandas） #

import pandas as pd

# 构造示例数据
data = {
    'Gene': ['MAD1L1', 'OAS1', 'OAS1'],
    'Drug': ['Lithium[17]', 'Lithium[17]', 'Lithium[7]'],
    'ID': ['34718328', '34718328', '27401222'],
    'Cite': ['[17]', '[17]', '[7]']
}
df = pd.DataFrame(data, index=[1,2,3])

# 提取药物基础名称（去除引用后缀），用于分组
df['Drug_Base'] = df['Drug'].str.extract(r'([a-zA-Z]+)')

# 分组聚合：合并药物引用，保留首个ID和Cite（可按需调整Cite合并逻辑）
merged_df = df.groupby(['Gene', 'Drug_Base']).agg(
    Drug=('Drug', lambda x: ''.join(x)),
    ID=('ID', 'first'),
    Cite=('Cite', 'first')
).reset_index(drop=True)

print(merged_df)

代码说明 #

1. 提取药物基础名称：通过正则表达式提取药物核心名称（如Lithium），确保同一药物的不同引用行能被归为一组。
2. 分组聚合：按Gene和Drug_Base分组，对各列做如下处理：
   • Drug：拼接所有带引用的药物字符串，得到合并后的完整药物名。
   • ID：取分组内首个ID值（与示例输出一致，若需合并所有ID可修改逻辑）。
   • Cite：示例中取首个Cite，若需合并所有引用，可替换为：

     Cite=('Cite', lambda x: '[' + ','.join([c.strip('[]') for c in x]) + ']')
     

     合并后Cite会变为[17,7]。

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