Getting mask of numpy array according to vector-type elements

Question

We can find indices of a scalar in numpy array like below:

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

When element is a vector:

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]]

I actually want to generate a similar mask like in the first code snippet in the second example.

mask = [False,True,False]

Is this possible in numpy library?

user2285236 · Accepted Answer · 2017-01-01 13:40:43Z

1

Since the comparison is element-wise, you need to reduce it using all on the first axis:

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

answered Jan 1, 2017 at 13:40

user2285236

Sign up to request clarification or add additional context in comments.

1 Comment

Thank you for the hint @ayhan