Python 2.7+Numpy中坐标列表的半径邻域非零值高效查找

嘿，针对你这个需求，我来分享几个高效的实现思路，尽量避免Python层面的循环，利用NumPy的向量运算或者可选的SciPy工具来提速——毕竟哪怕现在只有100个点，高效的方法扩展性也更好～

第一步：预处理非零坐标（必做优化） #

首先我们可以先过滤出所有values非零对应的坐标，这样后续只需要和这些点计算距离，减少不必要的计算量：

import numpy as np

# 假设你的coords和values已经填充好数据
coords = np.zeros((100, 2), dtype=np.int)
values = np.zeros(100, dtype=np.int)

# 提取非零值对应的坐标和数值
non_zero_mask = values != 0
non_zero_coords = coords[non_zero_mask]
non_zero_values = values[non_zero_mask]

r = 5  # 你的搜索半径

方法一：纯NumPy广播实现（无额外依赖） #

利用NumPy的广播机制，一次性计算所有坐标到非零坐标的距离平方（省去开根号的开销，速度更快），然后筛选出符合半径条件的邻点：

# 计算每个坐标与所有非零坐标的差值：形状为(100, N, 2)，N是非零点数量
diff = coords[:, np.newaxis, :] - non_zero_coords[np.newaxis, :, :]
# 计算距离平方（欧氏距离），形状为(100, N)
dist_sq = np.sum(diff ** 2, axis=2)

# 找到距离平方 <= r²的邻点掩码
neighbor_mask = dist_sq <= r ** 2

# 遍历每个坐标获取邻点
for i in range(len(coords)):
    # 获取当前坐标的邻点索引
    idx = np.where(neighbor_mask[i])[0]
    if len(idx) > 0:
        current_neighbors = non_zero_coords[idx]
        current_vals = non_zero_values[idx]
        print "坐标{}的邻点：{}，对应值：{}".format(coords[i], current_neighbors, current_vals)

优势： #

• 完全依赖NumPy，不需要额外安装库；
• 底层是C实现的向量运算，比Python循环快得多；
• 整数运算无精度损失，适合你的坐标类型。

方法二：SciPy KDTree加速（适合大数据量） #

如果你的数据量以后可能扩大，或者非零点数量很多，可以用SciPy的KDTree来做范围查询，查询复杂度是O(logN)，比全量计算更高效：

from scipy.spatial import KDTree

# 构建非零坐标的KDTree索引
kdtree = KDTree(non_zero_coords)

# 查询每个坐标在半径r内的邻点（返回每个点对应的邻点在non_zero_coords中的索引列表）
neighbor_idx_lists = kdtree.query_ball_point(coords, r=r)

# 遍历处理每个坐标的邻点
for i, idx_list in enumerate(neighbor_idx_lists):
    if idx_list:
        current_neighbors = non_zero_coords[idx_list]
        current_vals = non_zero_values[idx_list]
        print "坐标{}的邻点：{}，对应值：{}".format(coords[i], current_neighbors, current_vals)

注意： #

• 需要确保你的Python 2.7环境安装了SciPy（pip install scipy==1.2.3，这是支持Python 2.7的最后一个SciPy版本）；
• 当非零点数量远大于100时，KDTree的优势会非常明显。

选择建议 #

• 如果非零点数量很少（比如个位数），纯NumPy广播的方法足够快，代码也更简洁；
• 如果非零点数量较多，或者未来数据量会增长，优先选KDTree方案。

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