如何通过单次np.where调用获取数组的零元素与非零元素索引

你可以通过只做一次数组比较生成布尔掩码，再分别从掩码和取反掩码里提取索引，既减少重复计算，又避开低效的集合操作：

import numpy as np

# 仅执行一次数组比较，生成零元素的布尔掩码
mask = A == 0
zero_indices = np.nonzero(mask)
nonzero_indices = np.nonzero(~mask)

为什么这个方案更快？ #

• 你原来的两次np.where会重复执行A==0和A!=0两次数组比较，而上面的代码只做一次比较生成掩码，后续的索引提取是基于现成的布尔数组，开销比重复计算小很多。
• 你之前用的np.setdiff1d需要对整个数组的索引序列做排序、集合差运算，大型数组下这类操作的额外计算量极大，自然速度骤降；而np.nonzero是直接对布尔数组做向量式索引提取，完全利用numpy的底层优化，性能接近原生数组操作。

习惯用np.where的等价写法 #

要是更顺手用np.where，也可以基于同一掩码实现，效果和np.nonzero完全一致：

mask = A == 0
zero_indices = np.where(mask)
nonzero_indices = np.where(~mask)

实测下来，这种方式的速度和两次np.where相当甚至略快（少了一次数组比较），同时满足你“尽量减少重复数组操作”的需求。

内容的提问来源于stack exchange，提问作者C.M.O.B.