基于字典值生成新FASTA文件的For循环优化问询

问题背景 #

现有Python代码可实现需求：根据字典中键对应的登录号列表，从FASTA文件中筛选对应序列，生成以键命名的新FASTA文件。但处理50000条序列时速度偏慢，需要优化。

原代码存在的核心问题：

• 嵌套循环层级过多：遍历字典键→遍历所有FASTA行→遍历每个键的登录号，时间复杂度为O(KNM)，数据量大时效率极低
• 一次性加载全部FASTA行到内存，占用内存资源高
• 用j in line做字符串匹配，容易出现误匹配（比如登录号是其他编号的子串）
• 重复打开输出文件，IO操作冗余

优化思路 #

1. 建立登录号到输出文件的直接映射：提前把所有登录号和对应的输出文件对象绑定，避免多次查找匹配
2. 流式处理FASTA文件：逐行读取，不一次性加载全部内容，降低内存消耗
3. 精准匹配登录号：从FASTA标题行中提取完整登录号，避免子串误匹配

优化后的代码 #

# 预处理：建立登录号到输出文件对象的映射
acc_to_file = {}
output_files = []  # 保存所有打开的文件对象，最后统一关闭

for group_name, acc_list in d.items():
    f = open(group_name, 'w')
    output_files.append(f)
    for acc in acc_list:
        acc_to_file[acc.strip()] = f  # 去除登录号可能的空格，确保匹配准确

# 流式遍历FASTA文件，处理序列
current_header_acc = None
current_seq = []

with open('file_name.fasta', 'r') as seqs:
    for line in seqs:
        line = line.strip()
        if not line:
            continue  # 跳过空行
        if line.startswith('>'):
            # 先处理上一条序列（如果有的话）
            if current_header_acc is not None and current_header_acc in acc_to_file:
                acc_to_file[current_header_acc].write(f">{current_header_acc}\n{''.join(current_seq)}\n")
            # 提取当前登录号：从标题行取第一个字段，忽略后续注释内容
            current_header_acc = line.lstrip('>').strip().split()[0]
            current_seq = []
        else:
            # 累加序列行（支持多行序列的FASTA格式）
            current_seq.append(line)
    # 处理文件末尾的最后一条序列
    if current_header_acc is not None and current_header_acc in acc_to_file:
        acc_to_file[current_header_acc].write(f">{current_header_acc}\n{''.join(current_seq)}\n")

# 统一关闭所有输出文件，避免资源泄漏
for f in output_files:
    f.close()

优化效果说明 #

• 时间复杂度降至O(N)：仅遍历FASTA文件一次，每条序列的匹配写入操作是O(1)的字典查找
• 内存占用大幅降低：无需加载全部FASTA内容到内存，仅保存当前处理的序列片段
• 匹配更精准：提取标题行中的完整登录号，避免子串误匹配的情况
• IO操作更高效：提前打开所有输出文件，避免重复打开/关闭的开销

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