Python：检查DataFrame逗号分隔值匹配并计算样本均值

现有两个规模较大的DataFrame，数据如下：

df1：

gene_callers_id
0      4717766,4743899,11597717,12116240
1      4717766,4743899,12116240,7719716,4022000
2      4717766,4743899,12116240,7248697,7719716

df2：

gene_callers_id  sample_1  sample_2  
0               4743899  0.345000  0.176000  
1               4717766  0.000000  2.500000  
2               4743898  0.000000  0.684982 

需求：检查df1每行中逗号分隔的gene_callers_id值是否存在于df2的gene_callers_id列中，匹配则填入df2对应的sample_1、sample_2值，无匹配则填0；同时计算每行样本值的平均值，最终输出格式如下：

+------------------------------------------+-----------------+--------------+---------------------------+-----------+
| gene_callers_id                          | sample_1        | sample_1_avg |sample_2                   |sample2_avg|
+------------------------------------------+-----------------+--------------+---------------------------+-----------+
| 4717766,4743899,11597717,12116240        | 0,0.345000,0,0  | 0.08635      |2.500,0.684982,0,0         |0.79       |
+------------------------------------------+-----------------+--------------+---------------------------+-----------+
| 4717766,4743899,12116240,7719716,4022000 | 0,0,0,0,0       | 0            |2.500,0.684982,0,0,0       |0.64       |
+------------------------------------------+-----------------+--------------+---------------------------+-----------+
| 4717766,4743899,12116240,7248697,4743898 | 0,0,0,0,0.345000| 0.06900      |2.500,0.1760,0,0,0.684982  |0.67       |
+------------------------------------------+-----------------+--------------+---------------------------+-----------+

针对大规模数据集，优先用字典映射提升查找效率，避免循环中重复查询DataFrame，具体实现步骤如下：

1. 构建快速查找的映射字典

import pandas as pd

# 将df2转换为字典，key为gene_callers_id，value为对应的sample1和sample2值
sample_map = df2.set_index('gene_callers_id')[['sample_1', 'sample_2']].to_dict('index')

2. 定义每行数据的处理函数

def process_row(row):
    # 拆分逗号分隔的基因ID列表
    gene_ids = row['gene_callers_id'].split(',')
    sample1_vals = []
    sample2_vals = []
    
    for gid in gene_ids:
        # 统一转换为整数，匹配df2的ID类型
        gid_int = int(gid)
        # 查找映射值，无匹配则返回(0,0)
        vals = sample_map.get(gid_int, {'sample_1': 0, 'sample_2': 0})
        sample1_vals.append(str(vals['sample_1']))
        sample2_vals.append(str(vals['sample_2']))
    
    # 拼接成逗号分隔的字符串
    sample1_str = ','.join(sample1_vals)
    sample2_str = ','.join(sample2_vals)
    
    # 计算平均值，保留对应小数位数
    sample1_nums = [float(x) for x in sample1_vals]
    sample2_nums = [float(x) for x in sample2_vals]
    sample1_avg = round(sum(sample1_nums)/len(sample1_nums), 5) if sample1_nums else 0
    sample2_avg = round(sum(sample2_nums)/len(sample2_nums), 2) if sample2_nums else 0
    
    return pd.Series([sample1_str, sample1_avg, sample2_str, sample2_avg], 
                     index=['sample_1', 'sample_1_avg', 'sample_2', 'sample2_avg'])

3. 应用函数并合并结果

# 对df1每行应用处理函数，合并结果
result_df = df1.join(df1.apply(process_row, axis=1))

# 调整列顺序以匹配期望输出
result_df = result_df[['gene_callers_id', 'sample_1', 'sample_1_avg', 'sample_2', 'sample2_avg']]

# 打印最终结果
print(result_df.to_string(index=False))

关键说明 #

• 字典映射将查找时间复杂度从O(n*m)降至O(n)，大幅提升大规模数据的处理效率；
• 平均值的小数位数可通过调整round函数的参数自定义；
• 统一转换基因ID的类型，避免因类型不匹配导致的查找失败。

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