Python百万级列表构建字典及关联代码性能优化求助

原代码核心问题分析 #

你的代码运行缓慢的核心原因是两个关键步骤的时间复杂度过高：

1. 标签映射构建：原代码先收集所有标签，再对每个标签遍历所有短句子索引检查匹配，时间复杂度为O(M*K)（M为标签数，K为短句子数），百万级数据下会产生数亿次无效操作。
2. 正样本对生成：通过双重循环枚举所有i<j的句子对并检查集合交集，时间复杂度为O(K²)，哪怕K是100万，总操作数会达到5e11量级，完全无法完成计算。

优化后的代码实现 #

from collections import defaultdict
from itertools import combinations

# 1. 筛选短句子并收集有效标签集合（同步完成，避免重复遍历原始列表）
short_sent_indices = []
label_sets = []
max_length = 你的最大长度阈值

for idx, sent in enumerate(A_flat):
    # 筛选符合长度要求的句子
    if len(sent.split()) <= max_length:
        # 过滤掉'None'标签，生成有效标签集合
        valid_labels = {lbl for lbl in B_flat[idx] if lbl != 'None'}
        # 跳过无有效标签的句子（无法形成正样本对）
        if valid_labels:
            short_sent_indices.append(idx)
            label_sets.append(valid_labels)

# 2. 构建标签到短句子内部索引的映射（O(K)复杂度，K为短句子数）
label_to_inner_indices = defaultdict(list)
for inner_idx, lbl_set in enumerate(label_sets):
    for lbl in lbl_set:
        label_to_inner_indices[lbl].append(inner_idx)

# 3. 生成正样本对（基于标签分组，避免全量枚举）
positive_pairs = set()
for indices in label_to_inner_indices.values():
    # 用itertools.combinations生成所有i<j的组合（C语言实现，速度远快于Python循环）
    positive_pairs.update(combinations(indices, 2))

# 若需要原始索引的正样本对，可执行以下转换
# positive_pairs_original = {(short_sent_indices[i], short_sent_indices[j]) for (i, j) in positive_pairs}

关键优化点说明 #

• 合并遍历操作：将短句子筛选和标签集合生成分步合并，仅遍历一次原始列表，减少不必要的IO开销。
• 反向构建标签映射：不再通过标签反向查找句子，而是遍历每个句子的标签并将索引添加到对应标签列表，时间复杂度从O(MK)降至O(K)（每个句子最多4个标签，总操作数约4K）。
• 基于标签分组生成正样本：利用“共享同一标签的句子互为正样本”的逻辑，仅对每个标签下的索引生成两两组合，时间复杂度从O(K²)降至O(S)（S为所有标签组内的两两对数之和，远小于K²）。
• 使用高效工具函数：itertools.combinations是C语言实现的工具，比Python手动双重循环速度提升数倍。
• 提前过滤无效数据：直接跳过无有效标签的句子，减少后续处理的数据量。

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