根据vector-type元素获取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)