大规模集合间元素匹配式相似度计算的Python代码性能优化求助

你的问题核心是处理20万个集合时的三重循环效率太低——原代码的时间复杂度是O(N²K)（N是集合数量，K是每个集合的平均元素数），这在N=2e5的规模下完全不可行。不过好在你提到唯一标签总数只有300个，这是关键突破口，我们可以利用这个特性把复杂度降到O(NL + L² + N*L)（L是标签总数，仅300），性能会提升几个数量级。

优化思路 #

核心是把集合间的交集计数问题，转化为向量点积计算：

1. 把每个集合转换成标签维度的二进制向量（某标签存在则对应位置为1，否则为0）
2. 两个集合的交集元素数量，就是它们对应向量的点积（对应位置都为1的数量之和）
3. 用优化过的稀疏矩阵乘法计算所有向量的点积，最后筛选每个集合的前X个高相似度集合

优化后的代码示例 #

import json
import numpy as np
from scipy.sparse import csr_matrix

# 替换为你的真实输入数据
input_data = {
    "list1": ["label1", "label2", "label3"],
    "list2": ["label2", "label3", "label4"],
    "list3": ["label3", "label4", "label5"],
    "list4": ["label4", "label5", "label6"]
}

# 1. 构建标签到索引的映射，统一标签维度
all_labels = list({label for lst in input_data.values() for label in lst})
label_to_idx = {label: i for i, label in enumerate(all_labels)}
num_labels = len(all_labels)
set_keys = list(input_data.keys())
num_sets = len(set_keys)

# 2. 构建稀疏矩阵（每行对应一个集合，列对应标签，非零值表示标签存在）
row_indices = []
col_indices = []
for set_idx, (key, labels) in enumerate(input_data.items()):
    for label in labels:
        row_indices.append(set_idx)
        col_indices.append(label_to_idx[label])
data = np.ones(len(row_indices), dtype=int)
sparse_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(num_sets, num_labels))

# 3. 计算相似度矩阵（点积结果就是交集元素数量）
# 稀疏矩阵乘法是高度优化的，远快于手动循环
similarity_matrix = sparse_matrix.dot(sparse_matrix.T).toarray()

# 4. 生成结果（这里演示全量输出，若只需要前100个可看下方优化）
output = {}
top_X = 100  # 你需要的前X个高相似度集合
for set_idx, key in enumerate(set_keys):
    # 快速筛选前top_X个高相似度集合（用np.argpartition比sorted更快）
    sims = similarity_matrix[set_idx]
    # 找到前top_X大的元素索引
    top_indices = np.argpartition(sims, -top_X)[-top_X:]
    # 按相似度从高到低排序
    top_indices_sorted = top_indices[np.argsort(-sims[top_indices])]
    # 构建结果字典
    top_sims = {set_keys[j]: sims[j] for j in top_indices_sorted}
    output[key] = top_sims

print(json.dumps(output, indent=2))

为什么这个方法更快？ #

• 原代码问题：三重循环带来4e10次操作（20万20万10），完全无法在合理时间内完成
• 优化后优势：
  1. 构建稀疏矩阵仅需2e6次操作（20万*10），几乎瞬间完成
  2. 稀疏矩阵乘法的复杂度约为1.2e8次操作（20万3002 + 300²），比原代码小了3个数量级
  3. 稀疏矩阵大幅节省内存：仅需存储2e6个非零元素，而全量矩阵需要4e10个元素（根本存不下）

额外优化建议 #

• 如果不需要全量相似度矩阵，可以在计算后直接丢弃，只保留每个集合的前100个结果，进一步节省内存
• 若你的集合元素有重复（不是严格集合），只需将向量从二进制改为计数向量即可，点积依然能正确计算共同元素的总数

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