使用numpy.where条件索引时，如何提升numpy大数组的索引速度？

针对你处理形状为(50000000, 3)的大型数组场景，原代码里的np.where是多余的性能开销点，直接用布尔掩码索引就能大幅提升速度，优化后的实现如下：

# 构建布尔掩码
mask = (array[:, 0] == value) | (array[:, 1] == value)
# 直接用掩码索引数组
x = array[mask]

优化原理 #

原代码中np.where会先把布尔掩码转换成索引数组，再通过索引提取元素，多了一层索引转换的中间步骤。而直接使用布尔掩码索引时，Numpy内部会直接基于掩码筛选元素，省略了索引转换的开销，对于超大数组来说，这个节省的时间非常明显。

额外优化建议 #

如果你的数组不是连续内存布局（比如经过多次切片、转置操作），可以先将数组转为连续内存：

array = np.ascontiguousarray(array)

连续内存的数组在进行掩码索引时，Numpy的内存访问效率会更高，能进一步提升筛选速度。

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