如何快速计算两个Pandas DataFrame交集样本数？（SNP与表型矩阵场景）

问题背景 #

现有两个矩阵：

• 个体×SNPs矩阵：snp_data，eid为个体ID，其余列对应SNP，值为1表示个体携带该SNP
• 个体×表型矩阵：pheno_matrix_df，eid为个体ID，其余列对应表型，值≥1表示个体患病
  需要生成SNP×表型矩阵，每个单元格代表同时携带对应SNP且患对应疾病的共同样本数量。原双重循环（529个SNP × 1502个表型）耗时超一天，需优化。

核心优化思路 #

双重循环的本质是重复计算样本交集，效率极低。利用矩阵乘法的向量化运算可实现批量计算：

• 将SNP矩阵转置为「SNP×个体」结构，表型矩阵保持「个体×表型」结构，两者相乘后，每个元素就是对应SNP和表型同时满足条件的样本数，完全替代循环中的交集计数逻辑。
• 底层依赖numpy的C级高效运算，速度比Python循环提升数个数量级。

优化后代码 #

import pandas as pd
import numpy as np

# 1. 对齐样本，确保两个矩阵的个体ID完全一致且顺序相同
common_eids = snp_data['eid'].intersection(pheno_matrix_df['eid'])
# 按统一索引筛选并排序，避免计算时索引不匹配
snp_filtered = snp_data[snp_data['eid'].isin(common_eids)].set_index('eid').sort_index()
pheno_filtered = pheno_matrix_df[pheno_matrix_df['eid'].isin(common_eids)].set_index('eid').sort_index()

# 2. 转换为二进制矩阵：SNP取1的样本标记为True，表型≥1的样本标记为True
snp_binary = snp_filtered == 1
pheno_binary = pheno_filtered >= 1

# 3. 矩阵乘法批量计算共同样本数：SNP×个体  ×  个体×表型 = SNP×表型
result_matrix = np.dot(snp_binary.T.values, pheno_binary.values)

# 4. 转换为目标格式的DataFrame
snp_phenotype_matrix = pd.DataFrame(
    result_matrix,
    index=snp_binary.columns,
    columns=pheno_binary.columns
)

# 保存结果
snp_phenotype_matrix.to_csv("sample_count_matrix.csv")

关键优化点说明 #

• 样本对齐：通过set_index和sort_index确保两个矩阵的个体顺序完全一致，避免索引混乱导致的计算错误
• 向量化运算：矩阵乘法一次性完成所有SNP与表型的组合计算，彻底消除Python循环的性能开销
• 底层优化：用numpy数组执行计算，比纯pandas操作更高效，充分利用C语言的运算优势

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