请求优化坐标点匹配矩形区域zone_id的Python代码

我现在碰到个性能瓶颈：需要给df1里的每一个经纬度坐标点，匹配它所属矩形区域的zone_id。其中df1存储的是带经纬度的点数据，df2则是用左下角（lat_bl、lon_bl）和右上角（lat_tr、lon_tr）坐标定义的矩形区域，每个区域对应一个zone_id（我用df2的索引来设置的）。

但目前写的代码运行起来特别慢，尤其是数据量上去之后完全扛不住，想请教各位大佬怎么优化提速？

我当前的低效代码： #

def zone_map(df1, df2):
    df2['zone_id'] = df2.index
    for t ,t2 in df2.iterrows():
        mask=(df1['lat'] >=df2.loc[t,'lat_bl']) & (df1['lat'] <df2.loc[t,'lat_tr']) & (df1['lon'] >=df2.loc[t,'lon_bl']) & (df1['lon'] <df2.loc[t,'lon_tr'])
        df1.loc[mask, 'zone_id'] = t2['zone_id']
    return df1

问题根源分析 #

原代码用iterrows()遍历每个区域，每次遍历都要对整个df1做布尔索引判断，时间复杂度是O(N*M)（N是df2的区域数，M是df1的点数）。当数据量达到万级甚至十万级时，这种嵌套遍历的开销会直接爆炸。

优化方案推荐 #

方案1：向量化+merge_asof（适合有序不重叠区域）

如果你的矩形区域在纬度（或经度）上是不重叠且按区间有序排列的，merge_asof是最优解，它能把时间复杂度降到O(N log N + M log M)：

import pandas as pd

def fast_zone_map(df1, df2):
    # 给区域设置zone_id
    df2['zone_id'] = df2.index
    # 对df1按纬度排序，df2按区域左下角纬度排序
    df1_sorted = df1.sort_values('lat').reset_index(drop=True)
    df2_sorted = df2.sort_values('lat_bl').reset_index(drop=True)
    
    # 用merge_asof匹配点所在的纬度区间
    merged = pd.merge_asof(df1_sorted, df2_sorted, left_on='lat', right_on='lat_bl', direction='backward')
    # 过滤掉纬度超出区域右上角的点
    merged = merged[merged['lat'] < merged['lat_tr']]
    # 再过滤经度区间
    merged = merged[(merged['lon'] >= merged['lon_bl']) & (merged['lon'] < merged['lon_tr'])]
    
    # 恢复原df1的索引和结构
    result = df1.merge(merged[['lat', 'lon', 'zone_id']], on=['lat', 'lon'], how='left')
    return result

方案2：空间索引（适合任意区域分布，大数据量）

如果你的区域是任意分布、可能重叠的，推荐用geopandas构建空间索引，空间查找的效率会比遍历高几个数量级：

import geopandas as gpd
from shapely.geometry import Polygon

def geopandas_zone_map(df1, df2):
    df2['zone_id'] = df2.index
    # 把每个矩形区域转换成Shapely Polygon对象
    df2['geometry'] = df2.apply(lambda row: Polygon([
        (row['lon_bl'], row['lat_bl']),
        (row['lon_tr'], row['lat_bl']),
        (row['lon_tr'], row['lat_tr']),
        (row['lon_bl'], row['lat_tr']),
        (row['lon_bl'], row['lat_bl'])
    ]), axis=1)
    # 构建GeoDataFrame并创建空间索引
    gdf_zones = gpd.GeoDataFrame(df2, crs="EPSG:4326")
    _ = gdf_zones.sindex  # 显式创建空间索引
    
    # 把df1的坐标转换成点对象，构建GeoDataFrame
    df1['geometry'] = gpd.points_from_xy(df1['lon'], df1['lat'])
    gdf_points = gpd.GeoDataFrame(df1, crs="EPSG:4326")
    
    # 空间连接：找到每个点所在的区域
    joined = gpd.sjoin(gdf_points, gdf_zones, predicate='within', how='left')
    # 保留原df1的列和匹配到的zone_id
    result = joined[df1.columns.tolist() + ['zone_id']]
    return result

方案3：Numpy广播（适合中小数据量）

如果数据量不算特别大（比如点数和区域数都在万级以内），可以用Numpy的广播特性一次性计算所有匹配关系，避免循环：

import numpy as np
import pandas as pd

def numpy_zone_map(df1, df2):
    df2['zone_id'] = df2.index
    # 把数据转换成Numpy数组，方便广播计算
    lat_points = df1['lat'].values[:, np.newaxis]
    lon_points = df1['lon'].values[:, np.newaxis]
    lat_bl = df2['lat_bl'].values
    lat_tr = df2['lat_tr'].values
    lon_bl = df2['lon_bl'].values
    lon_tr = df2['lon_tr'].values
    zone_ids = df2['zone_id'].values
    
    # 广播计算所有点是否在每个区域内
    lat_mask = (lat_points >= lat_bl) & (lat_points < lat_tr)
    lon_mask = (lon_points >= lon_bl) & (lon_points < lon_tr)
    total_mask = lat_mask & lon_mask
    
    # 找到每个点匹配的第一个zone_id（若有多个重叠区域，取第一个；无匹配则设为NaN）
    matched_indices = np.argmax(total_mask, axis=1)
    df1['zone_id'] = np.where(total_mask.any(axis=1), zone_ids[matched_indices], np.nan)
    return df1

方案选择建议 #

• 区域有序不重叠 → 选方案1，速度最快，内存开销小
• 区域任意分布、数据量大 → 选方案2，空间索引的查找效率碾压遍历
• 中小数据量 → 选方案3，代码简洁易维护，比原循环快10~100倍

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