Python中高效比对坐标字典值是否在给定阈值范围内的方法求助

优化思路 #

原有代码的性能瓶颈主要来自4个方面：

• 手动写双层循环生成key对，且嵌套遍历坐标时重复调用list.index()做O(M)复杂度的索引查找，冗余开销大
• 用集合对称差计算坐标差值的逻辑存在漏洞（当两个坐标数组差异元素超过2个时会直接索引报错），且集合转换开销高
• 纯Python原生循环做差值计算，没有利用向量化运算能力
• 阈值difference用全局变量传入，耦合性高

优化实现 #

import itertools
import numpy as np

def CompareCoordinates(coordinate_dict, threshold):
    error_count = 0
    error_actors = {}
    key_list = sorted(coordinate_dict.keys())
    
    # 预转换所有坐标为numpy数组，同时保留1开头的计数索引
    preprocessed = {}
    for key in key_list:
        coords = []
        for idx, coord in enumerate(coordinate_dict[key], 1):
            coords.append({
                'idx': idx,
                'x': np.array(coord['x']),
                'y': np.array(coord['y']),
                'z': np.array(coord['z'])
            })
        preprocessed[key] = coords
    
    # 生成不重复的key对，代替手动双层循环
    for k1, k2 in itertools.combinations(key_list, 2):
        # 初始化对应存储键
        overlap_key = f'{k1} - {k2} - OVERLAP'
        x_key = f'{k1} - {k2} - X'
        y_key = f'{k1} - {k2} - Y'
        z_key = f'{k1} - {k2} - Z'
        error_actors.update({k: [] for k in [overlap_key, x_key, y_key, z_key]})
        
        # 遍历两组坐标
        for coord1 in preprocessed[k1]:
            idx1 = coord1['idx']
            x1, y1, z1 = coord1['x'], coord1['y'], coord1['z']
            for coord2 in preprocessed[k2]:
                idx2 = coord2['idx']
                x2, y2, z2 = coord2['x'], coord2['y'], coord2['z']
                
                # 先判断完全重叠
                if np.array_equal(x1, x2) and np.array_equal(y1, y2) and np.array_equal(z1, z2):
                    error_count += 1
                    error_actors[overlap_key].append([idx1, idx2])
                    continue
                
                # 逐个轴计算最大差值，判断是否符合阈值要求
                max_diff_x = np.max(np.abs(x1 - x2))
                if max_diff_x <= threshold:
                    error_count += 1
                    error_actors[x_key].append([f'{idx1}. coordinate', f'{idx2}. coordinate'])
                
                max_diff_y = np.max(np.abs(y1 - y2))
                if max_diff_y <= threshold:
                    error_count += 1
                    error_actors[y_key].append([f'{idx1}. coordinate', f'{idx2}. coordinate'])
                
                max_diff_z = np.max(np.abs(z1 - z2))
                if max_diff_z <= threshold:
                    error_count += 1
                    error_actors[z_key].append([f'{idx1}. coordinate', f'{idx2}. coordinate'])
    
    # 过滤空结果
    error_actors = {k: v for k, v in error_actors.items() if v}
    return error_count, error_actors

额外优化建议 #

如果数据量特别大，还可以做进一步优化：

• 对坐标值做哈希分桶，只比对同一桶或相邻桶内的坐标，避免O(M²)的全量两两比对
• 使用多进程并行处理不同的key对，利用多核CPU性能
• 提前过滤完全不可能符合阈值要求的坐标对，减少无效计算

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