如何加速嵌套循环？百万级空间数据匹配优化求助

现有代码处理高维数据时性能瓶颈显著，具体代码及场景如下：

spatial_data = [[1,2,3,4,5,6], ..... , [1,2,3,4,50,60]] # 长度1000万
reference = [[9, 39, 22, 28, 25, 5], ...... , [5, 16, 12, 34, 3, 9]] # 长度100

for x in reference:
    result = [c for c in spatial_data if len(set(x).intersection(c)) <= 2]

当前单次遍历reference耗时约5分钟，需执行1000次，总耗时预计3.5天，常规写法优化无明显效果，以下是针对性提速方案：

核心优化方向：减少重复计算+利用数据特性 #

1. 预转换集合，避免重复开销 #

提前将所有数据转换为集合，避免循环内重复创建集合：

# 预处理：一次性转换所有元素为集合
spatial_sets = [set(item) for item in spatial_data]
reference_sets = [set(item) for item in reference]

# 优化后的循环
for x_set in reference_sets:
    result = [spatial_data[i] for i, s in enumerate(spatial_sets) if len(x_set & s) <= 2]

用&替代intersection，语法更简洁且性能略优。

2. 短路计数，提前终止判断 #

不计算完整交集长度，一旦交集元素超过2就停止遍历，减少无效计算：

def count_intersection_shortcut(set_a, set_b):
    count = 0
    for elem in set_a:
        if elem in set_b:
            count += 1
            if count > 2:
                return count
    return count

# 循环中使用短路计数
for x_set in reference_sets:
    result = [spatial_data[i] for i, s in enumerate(spatial_sets) if count_intersection_shortcut(x_set, s) <= 2]

该方法在交集元素较多的场景下，能大幅缩短单条判断耗时。

3. 并行化处理，利用多CPU核心 #

由于每个reference元素的处理逻辑独立，可通过多进程并行拆分任务：

from multiprocessing import Pool

def process_reference(x_set):
    return [spatial_data[i] for i, s in enumerate(spatial_sets) if count_intersection_shortcut(x_set, s) <= 2]

if __name__ == "__main__":
    spatial_sets = [set(item) for item in spatial_data]
    reference_sets = [set(item) for item in reference]
    
    # 根据CPU核心数调整进程数
    with Pool(processes=4) as pool:
        results = pool.map(process_reference, reference_sets)

并行化可将总耗时压缩至接近单核心耗时除以核心数（忽略进程通信开销）。

4. 向量化运算（数值型数据专属） #

若所有元素为整数，可借助NumPy的向量化运算加速：

import numpy as np

# 转换为NumPy二维数组
spatial_np = np.array(spatial_data)
reference_np = np.array(reference)

for x in reference_np:
    # 向量化计算每个spatial元素与x的交集数量
    intersection_counts = np.sum(np.isin(spatial_np, x), axis=1)
    # 筛选符合条件的元素
    result = spatial_data[intersection_counts <= 2]

NumPy底层基于C实现，比纯Python循环效率提升数倍。

5. 倒排索引，空间换时间 #

构建倒排索引记录每个数值对应的spatial元素位置，大幅减少需要判断的元素范围：

from collections import defaultdict

# 构建倒排索引：key为数值，value为包含该数值的spatial元素索引列表
inverted_index = defaultdict(list)
for idx, item in enumerate(spatial_data):
    for num in item:
        inverted_index[num].append(idx)

# 处理每个reference元素
for x in reference:
    # 获取所有与x有交集的候选元素索引
    candidate_indices = set()
    for num in x:
        candidate_indices.update(inverted_index.get(num, []))
    # 对候选元素做精确判断
    result = []
    for idx in candidate_indices:
        if len(set(x) & spatial_sets[idx]) <= 2:
            result.append(spatial_data[idx])
    # 补充完全无交集的元素（交集数为0，符合<=2条件）
    total_indices = set(range(len(spatial_data)))
    no_intersection_indices = total_indices - candidate_indices
    result.extend([spatial_data[idx] for idx in no_intersection_indices])

该方案适合reference元素数值在spatial_data中出现频率较低的场景。

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