向量化经纬度与网格匹配循环 提升大规模数据处理效率

问题背景 #

我拥有两个二维数组：

• locs_array：存储待匹配的经纬度点位，示例如下：

array([[-122.463425,   47.195741],
       [-122.498139,   47.190166]])

• parents_array：存储网格定义，列依次为[id, centerlon, centerlat, upperlon, lowerlon, upperlat, lowerlat]，示例如下：

array([[   1.        , -122.463425  ,   47.195741  , -122.46331271,
        -122.46353729,   47.19585367,   47.19562833],
       [   2.        , -122.498149  ,   47.190275  , -122.49803671,
        -122.49826129,   47.19038767,   47.19016233]])

当前实现是通过循环调用_findFirstParent函数逐个匹配点位，但locs_array包含超1亿条记录，parents_array包含超1百万条记录，循环效率极低，以下是针对性的优化方案：

优化方案 #

1. 全向量化运算（无Python循环） #

利用NumPy的广播机制，一次性完成所有点位的匹配判断，彻底规避Python循环的性能开销。

核心逻辑：扩展数组维度实现广播，生成布尔掩码矩阵快速筛选匹配网格，再提取对应ID。

import numpy as np

# 提取parents_array中的关键列
upper_lon = parents_array[:, 3]
lower_lon = parents_array[:, 4]
upper_lat = parents_array[:, 5]
lower_lat = parents_array[:, 6]
parent_ids = parents_array[:, 0]

# 扩展维度实现广播：locs_array(N,2) → (N,1,2)，网格参数(M,) → (1,M)
lons = locs_array[:, 0, np.newaxis]
lats = locs_array[:, 1, np.newaxis]

# 生成匹配掩码
lon_mask = (upper_lon >= lons) & (lower_lon <= lons)
lat_mask = (upper_lat >= lats) & (lower_lat <= lats)
total_mask = lon_mask & lat_mask

# 提取每个点位的第一个匹配网格ID，无匹配则设为-1
matches = np.full(len(locs_array), -1, dtype=int)
valid_indices = total_mask.any(axis=1)
matches[valid_indices] = parent_ids[total_mask[valid_indices].argmax(axis=1)]

注意：若1亿×1百万的掩码矩阵超出内存，可将locs_array拆分为小批次（如每批10万条）逐批处理。

2. 空间索引加速（KDTree/R-tree） #

通过空间索引快速缩小候选网格范围，避免遍历全部网格。

方案A：KDTree（基于网格中心快速定位）

from scipy.spatial import KDTree

# 用网格中心坐标构建KDTree
tree = KDTree(parents_array[:, 1:3])

# 每个点位先查询最近的k个网格（k建议取5-10），再在候选中验证匹配
k = 5
distances, candidate_indices = tree.query(locs_array, k=k)

matches = np.full(len(locs_array), -1, dtype=int)
for i in range(len(locs_array)):
    for idx in candidate_indices[i]:
        if (parents_array[idx,3] >= locs_array[i,0] 
            and parents_array[idx,4] <= locs_array[i,0]
            and parents_array[idx,5] >= locs_array[i,1]
            and parents_array[idx,6] <= locs_array[i,1]):
            matches[i] = parents_array[idx,0]
            break

方案B：R-tree（基于网格边界的精确查询）

适合复杂空间范围的快速匹配，直接查询包含目标点位的网格：

from rtree import index

# 创建R-tree索引，存储网格边界与ID
idx = index.Index()
for i in range(len(parents_array)):
    # R-tree边界格式：(min_lon, min_lat, max_lon, max_lat)
    bounds = (parents_array[i,4], parents_array[i,6], parents_array[i,3], parents_array[i,5])
    idx.insert(i, bounds, obj=parents_array[i,0])

# 批量查询每个点位对应的网格
matches = []
for lon, lat in locs_array:
    result = list(idx.intersection((lon, lat, lon, lat), objects=True))
    matches.append(result[0].object if result else -1)
matches = np.array(matches)

3. 网格坐标离散化（仅适用于规则网格） #

若parents_array是等间隔规则网格，可直接通过计算离散坐标匹配网格，时间复杂度O(N)，速度最快。

# 计算网格的经纬度间隔（假设规则）
dlon = np.unique(np.diff(np.sort(parents_array[:,1])))[0]
dlat = np.unique(np.diff(np.sort(parents_array[:,2])))[0]
lon0 = np.min(parents_array[:,1]) - dlon/2
lat0 = np.min(parents_array[:,2]) - dlat/2

# 计算每个点位对应的网格索引
lon_idx = np.floor((locs_array[:,0] - lon0) / dlon).astype(int)
lat_idx = np.floor((locs_array[:,1] - lat0) / dlat).astype(int)

# 根据网格ID的生成规则计算匹配ID（示例：ID从1开始，按纬度×经度网格数+经度索引）
num_lon_grids = len(np.unique(parents_array[:,1]))
matches = lat_idx * num_lon_grids + lon_idx + 1

4. 多进程并行处理 #

结合多核CPU，将大批次数据拆分为小批次并行计算，进一步提升速度：

from multiprocessing import Pool

# 定义单批次处理函数（复用向量化逻辑）
def process_batch(batch):
    lons = batch[:,0, np.newaxis]
    lats = batch[:,1, np.newaxis]
    lon_mask = (upper_lon >= lons) & (lower_lon <= lons)
    lat_mask = (upper_lat >= lats) & (lower_lat <= lats)
    total_mask = lon_mask & lat_mask
    batch_matches = np.full(len(batch), -1, dtype=int)
    valid = total_mask.any(axis=1)
    batch_matches[valid] = parent_ids[total_mask[valid].argmax(axis=1)]
    return batch_matches

# 拆分批次
batch_size = 100000
batches = [locs_array[i:i+batch_size] for i in range(0, len(locs_array), batch_size)]

# 并行处理并合并结果
with Pool() as pool:
    results = pool.map(process_batch, batches)
matches = np.concatenate(results)

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