优化DataFrame坐标行距离匹配：替代双重循环的高效实现方案

原代码的双重循环时间复杂度为O(n²)，数据量稍大就会严重卡顿，以下是两种更高效的实现思路：

方案一：用scipy的KDTree实现球面距离查询 #

KDTree可快速完成最近邻搜索，结合haversine公式计算球面距离，适合大数据量场景。

步骤： #

• 先过滤经纬度缺失的行，或标记这些行的邻居为空
• 将经纬度转换为弧度（KDTree的haversine参数要求输入弧度）
• 构建KDTree，批量查询每个点在不同距离阈值内的邻居
• 映射回原DataFrame的id，生成结果列

代码示例： #

import pandas as pd
import numpy as np
from scipy.spatial import cKDTree

# 定义haversine距离计算（返回公里数）
def haversine(lon1, lat1, lon2, lat2):
    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    r = 6371  # 地球半径（公里）
    return c * r

# 加载数据并处理缺失值
df = pd.read_csv("coordinates.csv", na_values=None)
valid_mask = df[['longitude', 'latitude']].notna().all(axis=1)
valid_coords = df.loc[valid_mask, ['latitude', 'longitude']].values
valid_coords_rad = np.radians(valid_coords)

# 构建KDTree
tree = cKDTree(valid_coords_rad, leafsize=10)

# 距离阈值转换为球面弧度差（KDTree查询用弧度）
thresholds = {
    '1km': 1 / 6371,
    '5km': 5 / 6371,
    '10km': 10 / 6371,
    '50km': 50 / 6371
}

# 初始化结果列
for col in thresholds.keys():
    df[col] = [[] for _ in range(len(df))]

# 批量查询有效坐标的邻居
valid_indices = df[valid_mask].index.tolist()
for idx in valid_indices:
    point_idx_in_valid = valid_indices.index(idx)
    point = valid_coords_rad[point_idx_in_valid]
    for dist_name, rad_dist in thresholds.items():
        neighbor_positions = tree.query_ball_point(point, rad_dist)
        # 排除自身
        neighbor_positions = [p for p in neighbor_positions if p != point_idx_in_valid]
        # 映射回原id
        neighbor_ids = df.loc[valid_indices].iloc[neighbor_positions]['id'].tolist()
        df.at[idx, dist_name] = neighbor_ids

方案二：用Geopandas空间索引（直观地理数据处理） #

如果熟悉地理数据操作，Geopandas的空间索引可直接用公里作为阈值查询，代码更易读。

步骤： #

• 将DataFrame转换为GeoDataFrame，创建点几何对象
• 转换为投影坐标系（支持米为单位，减少距离计算误差）
• 构建空间索引，批量查询每个点在不同距离内的邻居
• 提取邻居id并生成结果列

代码示例： #

import pandas as pd
import geopandas as gpd
from shapely.geometry import Point

# 加载数据
df = pd.read_csv("coordinates.csv", na_values=None)

# 转换为GeoDataFrame（WGS84坐标系）
gdf = gpd.GeoDataFrame(
    df,
    geometry=gpd.points_from_xy(df.longitude, df.latitude),
    crs="EPSG:4326"
)
# 转换为UTM投影（示例用北京附近的EPSG:32633，需根据数据区域调整）
gdf = gdf.to_crs("EPSG:32633")

# 构建空间索引
sindex = gdf.sindex

# 距离阈值（单位：米）
thresholds = {
    '1km': 1000,
    '5km': 5000,
    '10km': 10000,
    '50km': 50000
}

# 初始化结果列
for col in thresholds.keys():
    gdf[col] = [[] for _ in range(len(gdf))]

# 遍历查询邻居
for idx, row in gdf.iterrows():
    if pd.isna(row.geometry):
        continue
    # 生成扩大后的边界框，缩小候选范围
    bbox = row.geometry.buffer(thresholds['50km']).bounds
    candidate_indices = list(sindex.intersection(bbox))
    # 排除自身
    candidate_indices = [i for i in candidate_indices if i != idx]
    candidates = gdf.iloc[candidate_indices]
    # 计算距离
    distances = candidates.geometry.distance(row.geometry)
    # 按阈值筛选id
    for dist_name, dist_m in thresholds.items():
        neighbor_ids = candidates[distances < dist_m]['id'].tolist()
        gdf.at[idx, dist_name] = neighbor_ids

# 转回普通DataFrame（可选）
df = pd.DataFrame(gdf.drop(columns='geometry'))

方案对比 #

• KDTree方案：速度更快，适合10万+级别的超大数据量，需处理弧度转换，对地理坐标逻辑要求稍高
• Geopandas方案：代码直观，支持复杂地理操作，适合中等数据量，需注意选择对应区域的投影坐标系

额外优化建议 #

• 提前过滤缺失值：避免循环中重复判断空值，减少无效计算
• 避免重复计算：KDTree和空间索引会自动跳过(i,j)与(j,i)的重复距离计算
• 批量操作优先：尽量用numpy/Geopandas的向量化方法代替Python循环，降低开销

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