python - 使用 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<b>[z % 3 == 0]</b>

对于您的示例矩阵，这将生成:

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

如果您传递一个维度与 bool 值相同的矩阵作为索引，您将得到一个包含该矩阵元素的数组，其中 bool 矩阵为 True。

这将进一步提高工作效率，因为您在 numpy 级别进行过滤(而列表理解在 Python 解释器级别工作)。