Using np.where to find element from sub-arrays

user_9853078 0
Viewed 1240

I'm trying to find all elements in array X that match element on array Y using np.where() and the condition on where() function is comparing list (a) not one element. Please see the following code:

X = np.array([[0, 2], [2, 1], [1, 3], [5, 9], [6, 7], [4, 6]])
Y = np.array([1, 2, 3, 4, 4, 5])
a = [2, 3, 4]
matchedX = X[np.where(Y == a)]

I am expecting to get the result like this:

array([[2, 1],
   [1, 3],
   [5, 9],
   [6, 7]])

but I got different results:

array([], shape=(0, 2), dtype=int64)

So, I need an alternative solution in a way that I can get the same elements if I do not know about the values of a? This line below give me the exact results that I want, but I do not know in previous the a values.

matchedX = X[np.where((Y == 2) | (Y==3) | (Y==4))]
numpy python
2 Answers

You can use the set functions of numpy:

X[np.where(np.isin(Y, a))]

array([[2, 1],
       [1, 3],
       [5, 9],
       [6, 7]])
user_42674140

You can skip the np.where, which is redundant here, and just index using np.isin:

X[np.isin(Y,a)]

array([[2, 1],
       [1, 3],
       [5, 9],
       [6, 7]])

This is because np.isin gives you a boolean array of where Y is in a:

array([False,  True,  True,  True,  True, False])

So by indexing with this array, it only selects those rows where True

user_66711760
Related