python - 根据向量类型元素获取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)