根据向量类型元素获取numpy数组的掩码

Question

我们可以在 numpy 数组中找到标量的索引，如下所示：

import numpy as np
array = np.array([1,2,3])
mask = (array == 2) #mask is now [False,True,False]

当元素是向量时：

import numpy as np
array = np.array([[1,2],[1,4],[5,6]])
mask = (array == [1,4]) #mask is now [[True,False],[True,True],[False,False]]

我实际上想生成与第二个示例中的第一个代码片段类似的掩码。

mask = [False,True,False]

这在 numpy 库中可行吗？

Answer 1

由于比较是按元素进行的，因此您需要在第一个轴上使用 all 来减少它：

(array == [1, 4]).all(axis=1)
Out: array([False,  True, False], dtype=bool)

Getting mask of numpy array according to vector-type elements