如何用NumPy快速筛选mask值为0的二维坐标数组？

你原来用map遍历每个点的方式属于Python层面的逐元素操作，数据量大的时候效率很低。NumPy的矢量化操作能解决这个问题，直接通过高级索引一次性完成判断和过滤：

优化后的代码 #

import numpy as np

mask = np.array([[0, 0, 0, 255, 255],
                 [0, 0, 0, 0, 255],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [255, 255, 0, 0, 0]])

pts0 = np.array([[1, 1], [1, 4], [0, 4], [2, 2], [3, 3]])

# 直接用矢量化索引获取mask对应位置的值，生成布尔掩码后过滤
pts0_filtered = pts0[mask[pts0[:, 1], pts0[:, 0]] == 0]

print(pts0_filtered)  # [[1 1]
                      #  [2 2]
                      #  [3 3]]

原理说明 #

• pts0[:, 1]提取所有点的y坐标，pts0[:, 0]提取x坐标，NumPy支持用这两个一维数组作为索引，直接取出mask中对应位置的所有值（这一步是矢量化计算，没有Python循环）
• 用== 0生成布尔数组，标记哪些点符合条件
• 最后用布尔数组索引pts0，直接得到过滤后的结果

这种方式的计算全部在NumPy底层的C实现中完成，当处理大规模点集时，速度会比原来的map方法快很多，完全避免了Python循环的开销。

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