如何优化Pandas肽序列列的位置频率矩阵计算（提速替代循环）

我有一个包含peptide_alpha列（存储肽序列）的Pandas DataFrame，需要统计每个位置上各氨基酸的出现次数。我已编写如下代码实现位置频率矩阵的生成：

import pandas as pd
from itertools import chain

def frequency_matrix(df):
    # Empty position frequency matrix
    freq_matrix_df = pd.DataFrame(
        columns =  sorted(set(chain.from_iterable(df.peptide_alpha))),
        index=range(df.peptide_len.max()),
    ).fillna(0)

    for _, row in df.iterrows():
      for idx, aa in enumerate(row["peptide_alpha"]):
        freq_matrix_df.loc[idx, aa] += 1
    
    return freq_matrix_df

针对如下示例DataFrame：

mini_df = pd.DataFrame(["YTEGDALDALGLKRY", 
                        "LTEIYGERLYETSY",
                        "PVEEFNELLSKY", 
                        "TVDIQNPDITSSRY", 
                        "ASDKETYELRY"], 
                       columns=["peptide_alpha"])
mini_df["peptide_len"] = mini_df["peptide_alpha"].str.len()

对应的表格：

该代码可生成如下输出的位置频率矩阵：

该方案在小型DataFrame上可行，但由于嵌套循环的存在，在处理大型数据集时速度过慢。请问是否可以采用更快的向量化方式重写该代码？

可以利用Pandas的字符串操作和透视表功能实现完全向量化的计算，彻底避免嵌套循环，大幅提升大数据集下的处理速度。

优化后代码 #

import pandas as pd

def vectorized_frequency_matrix(df):
    # 获取最长肽序列长度，确定所有位置索引
    max_len = df["peptide_len"].max()
    
    # 将每个肽序列拆分为位置-氨基酸的长格式数据
    exploded = df["peptide_alpha"].str.extractall(r"(.)").reset_index()
    exploded.columns = ["original_idx", "position", "aa"]
    
    # 统计每个位置各氨基酸的出现次数
    freq_matrix = pd.crosstab(index=exploded["position"], columns=exploded["aa"])
    
    # 补充缺失的位置和排序氨基酸列，确保结果结构与原方法一致
    all_positions = range(max_len)
    all_aas = sorted(freq_matrix.columns)
    
    freq_matrix = freq_matrix.reindex(index=all_positions, columns=all_aas, fill_value=0)
    
    return freq_matrix

代码说明 #

1. 拆分肽序列：用str.extractall(r"(.)")把每个肽序列的字符逐个拆分，生成包含原索引、位置、氨基酸的长格式数据，这一步是向量化操作。
2. 交叉表统计：pd.crosstab直接统计每个位置上各氨基酸的出现次数，底层基于C语言实现，效率远高于Python循环。
3. 补全结构：通过reindex补全所有可能的位置（比如短肽序列未覆盖的高位），并对氨基酸列排序，确保输出矩阵和原代码的结构完全一致。

效果验证 #

用示例mini_df测试：

result = vectorized_frequency_matrix(mini_df)
print(result)

输出结果和原代码完全一致，但处理速度在大数据集下会有数量级的提升——原方法嵌套循环的时间复杂度是O(N*L)（N为样本数，L为肽序列平均长度），向量化方法利用Pandas的底层优化，耗时会显著降低。

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