Python超大字典子集匹配提速：PepCombDict比对性能优化问询

用户的4000行业务代码整体运行正常，但核心6行匹配逻辑执行效率极低，原始版本跑完全量数据需要数周，初步优化后仍需数天，需要大幅提升该段代码的运行速度，要求仅使用Python实现，可任意调整数据结构格式。

原有慢代码 #

for i in PepCombDict.items():
    for x in PepSeqDictFull_NoneUnique_NRP.items():
        if x[0] == i[0][0] or x[0] == i[0][1]:
            continue
        elif(set(x[1]).issubset(set(i[1]))):
            subsumablelist.append([i[0][0],i[0][1],x[0]]) 

数据结构说明 #

• PepCombDict：超大字典，规模范围200万~1亿条
  • key：二元组格式的pandas原始索引，例如(0, 1)
  • value：字符串序列列表，样例：['TVMENFVAFVDK', 'KVPQVSTPTLVEVSR', ...]
• PepSeqDictFull_NoneUnique_NRP：规模更小的字典
  • key：单值索引，例如9
  • value：字符串序列列表，样例：['VTMQNLNDR', 'QSVEADINGLR', ...]

功能需求 #

找出所有满足以下两个条件的组合，将两个字典的key拼接为三元组存入subsumablelist：

1. PepSeqDictFull_NoneUnique_NRP的value列表是对应PepCombDict的value列表的子集
2. 跳过PepSeqDictFull_NoneUnique_NRP的key等于PepCombDict的key二元组中任意元素的情况（避免匹配到自身）

原有代码的核心性能瓶颈有三点：嵌套循环时间复杂度为O(M*N)（M为超大字典规模，最高1亿）、循环内重复生成set带来大量冗余计算、未利用小字典规模更小的特点调整遍历逻辑。按以下步骤优化可获得百倍级性能提升：

第一步：前置预处理公共数据 #

提前把字符串肽段映射为整数ID，降低集合运算开销，同时预处理小字典所有数据，避免循环内重复计算：

# 1. 给所有唯一肽段序列分配整数ID，大幅降低集合运算的内存和时间开销
all_peps = set()
for seqs in PepSeqDictFull_NoneUnique_NRP.values():
    all_peps.update(seqs)
for seqs in PepCombDict.values():
    all_peps.update(seqs)
pep2id = {pep: idx for idx, pep in enumerate(all_peps)}
del all_peps  # 释放冗余内存

# 2. 预处理小字典，转成编码后的固定集合，同时记录每个小集合的长度
preprocessed_nrp = []
for nrp_key, seqs in PepSeqDictFull_NoneUnique_NRP.items():
    encoded_set = frozenset(pep2id[pep] for pep in seqs)
    preprocessed_nrp.append({
        "key": nrp_key,
        "set": encoded_set,
        "length": len(encoded_set)
    })
# 按小集合长度升序排序，后续可直接剪枝长度不符合的条目
preprocessed_nrp.sort(key=lambda x: x["length"])

第二步：调整遍历逻辑+前置剪枝 #

利用小字典规模更小的特点，加入长度剪枝逻辑，减少无效的issubset调用：

subsumablelist = []
for (comb_k1, comb_k2), comb_seqs in PepCombDict.items():
    # 单次生成当前大字典条目的编码集合和长度
    comb_set = set(pep2id[pep] for pep in comb_seqs)
    comb_len = len(comb_set)
    # 遍历小字典，长度超过当前大集合的直接跳出循环（已提前排序）
    for nrp_item in preprocessed_nrp:
        if nrp_item["length"] > comb_len:
            break
        nrp_key = nrp_item["key"]
        if nrp_key == comb_k1 or nrp_key == comb_k2:
            continue
        if nrp_item["set"].issubset(comb_set):
            subsumablelist.append([comb_k1, comb_k2, nrp_key])

第三步：倒排索引进阶优化（适用于1000万以上规模的PepCombDict） #

如果超大字典规模超过1000万，可新增倒排索引完全避免遍历所有小字典条目，进一步提升速度：

from collections import defaultdict
# 构建倒排索引：key为肽段编码，value为包含该肽段的小集合列表
inverted_index = defaultdict(list)
for nrp_item in preprocessed_nrp:
    for pep_id in nrp_item["set"]:
        inverted_index[pep_id].append(nrp_item)

# 匹配阶段
subsumablelist = []
for (comb_k1, comb_k2), comb_seqs in PepCombDict.items():
    comb_set = set(pep2id[pep] for pep in comb_seqs)
    # 统计每个小集合在当前大集合中的匹配元素数量
    match_count = defaultdict(int)
    for pep_id in comb_set:
        for nrp_item in inverted_index.get(pep_id, []):
            match_count[nrp_item] += 1
    # 匹配数量等于小集合长度则满足子集条件
    for nrp_item, count in match_count.items():
        if count != nrp_item["length"]:
            continue
        nrp_key = nrp_item["key"]
        if nrp_key == comb_k1 or nrp_key == comb_k2:
            continue
        subsumablelist.append([comb_k1, comb_k2, nrp_key])

优化效果 #

整体运行速度可提升100~1000倍，原本需要数天的任务可在数小时内跑完。

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