如何向量化嵌套循环以优化成对距离计算脚本性能

核心思路 #

放弃iterrows()的Python层面嵌套循环，利用numpy广播机制一次性计算所有坐标点对的距离，再通过向量化筛选得到结果。底层由C实现的numpy操作能将速度提升几个数量级，彻底解决耗时过长的问题。

具体实现步骤 #

假设你的DataFrame包含平面坐标列（如x, y），若为经纬度坐标，见后续补充方案。

1. 提取坐标为numpy数组 #

将两个DataFrame的坐标列转换为numpy数组，方便广播运算：

import numpy as np
import pandas as pd

# 替换为你的实际坐标列名
q_coords = qinsy_file_2[['x', 'y']].to_numpy()
s_coords = segy_vlookup[['x', 'y']].to_numpy()

2. 向量化计算所有点对距离 #

利用numpy广播，一次性计算qinsy_file_2中每个点到segy_vlookup中所有点的欧氏距离：

# 计算坐标差的平方和，再开根号得到欧氏距离矩阵
# 最终得到形状为 (15000, 1500) 的距离矩阵，每个元素对应一对点的距离
distances = np.sqrt(((q_coords[:, np.newaxis, :] - s_coords[np.newaxis, :, :]) ** 2).sum(axis=2))

3. 筛选符合条件的行 #

根据需求选择以下两种筛选方式：

方式1：保留所有存在近邻点的qinsy行

如果只需要qinsy_file_2中至少有一个segy_vlookup点距离小于buffer的行：

buffer = 100  # 替换为你的预设距离阈值
# 检查每行是否存在距离小于buffer的点
has_close_point = (distances < buffer).any(axis=1)
out_df = qinsy_file_2[has_close_point].copy()

方式2：保留所有匹配的点对（含segy信息）

如果需要将符合条件的qinsy行和对应的segy行合并：

# 获取所有距离小于buffer的点对索引
q_indices, s_indices = np.where(distances < buffer)
# 合并两个DataFrame的对应行
out_df = pd.concat([
    qinsy_file_2.iloc[q_indices].reset_index(drop=True),
    segy_vlookup.iloc[s_indices].reset_index(drop=True)
], axis=1)

经纬度坐标适配方案 #

如果是地理经纬度坐标，不能用欧氏距离，需用Haversine公式计算球面距离，同样用向量化实现：

def haversine(lat1, lon1, lat2, lon2):
    """计算两点间的球面距离（单位：公里）"""
    R = 6371  # 地球平均半径（公里）
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    return R * c

# 提取经纬度数组
q_lat, q_lon = qinsy_file_2['lat'].to_numpy(), qinsy_file_2['lon'].to_numpy()
s_lat, s_lon = segy_vlookup['lat'].to_numpy(), segy_vlookup['lon'].to_numpy()

# 广播计算所有点对的球面距离
distances = haversine(q_lat[:, np.newaxis], q_lon[:, np.newaxis], s_lat[np.newaxis, :], s_lon[np.newaxis, :])

# 后续筛选步骤同平面坐标方案

性能说明 #

• 内存占用：15000×1500的距离矩阵约占180MB（float64类型），完全在普通机器的内存承受范围内。
• 速度对比：向量化操作将Python循环的O(n×m)时间复杂度转化为底层C级运算，耗时可从2小时压缩至数分钟甚至几十秒。

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