使用numpy.argwhere获取np.array中的匹配值

Question

我想np.argwhere()用来获取一个值np.array.

例如:

z = np.arange(9).reshape(3,3)

[[0 1 2]
 [3 4 5]
 [6 7 8]]

zi = np.argwhere(z % 3 == 0)

[[0 0]
 [1 0]
 [2 0]]

我想要这个数组:[0, 3, 6]并且这样做了:

t = [z[tuple(i)] for i in zi] # -> [0, 3, 6]

我认为有一种更简单的方法.

Answer 1

为什么不在这里简单地使用遮蔽:

z[z % 3 == 0]

对于您的样本矩阵,这将生成:

>>> z[z % 3 == 0]
array([0, 3, 6])

如果传递具有相同维度的矩阵作为索引的布尔值,则会得到一个数组,其中包含布尔矩阵所在的矩阵元素True.

这将更有效,因为您在numpy级别进行过滤(而列表理解在Python解释器级别工作).