Is there an array method for testing multiple equality values?

Question

I want to know where array a is equal to any of the values in array b.

For example,

a = np.random.randint(0,16, size=(3,4))
b = np.array([2,3,9])

# like this, but for any size b:
locations = np.nonzero((a==b[0]) | (a==b[1]) | (a==b[3]))

The reason is so I can change the values in a from (any of b) to another value:

a[locations] = 99

Or-ing the equality checks is not a great solution, because I would like to do this without knowing the size of b ahead of time. Is there an array solution?

[edit] There are now 2 good answers to this question, one using broadcasting with extra dimensions, and another using np.in1d. Both work for the specific case in this question. I ended up using np.isin instead, since it seems like it is more agnostic to the shapes of both a and b.
I accepted the answer that taught me about in1d since that led me to my preferred solution.

ALollz · Accepted Answer · 2021-02-21 23:14:25Z

1

You can use np.in1d then reshape back to a's shape so you can set the values in a to your special flag.

import numpy as np
np.random.seed(410012)

a = np.random.randint(0, 16, size=(3, 4))
#array([[ 8,  5,  5, 15],
#       [ 3, 13,  8, 10],
#       [ 3, 11,  0, 10]])

b = np.array([[2,3,9], [4,5,6]])

a[np.in1d(a, b).reshape(a.shape)] = 999
#array([[  8, 999, 999,  15],
#       [999,  13,   8,  10],
#       [999,  11,   0,  10]])

answered Feb 21, 2021 at 23:14

ALollz

59.7k7 gold badges73 silver badges97 bronze badges

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

1 Comment

Alex Shroyer Over a year ago

I didn't know about in1d, but it led me to isin, which I like even better. Thanks for teaching me a new thing about numpy.

Akshay Sehgal · Accepted Answer · 2021-02-21 21:39:48Z

1

Or-ing the equality checks is not a great solution, because I would like to do this without knowing the size of b ahead of time.

EDIT:

Vectorized equivalent to the code you have written above -

a = np.random.randint(0,16, size=(3,4))
b = np.array([2,3,9])

locations = np.nonzero((a==b[0]) | (a==b[1]) | (a==b[2]))

locations2 = np.nonzero((a[None,:,:]==b[:,None,None]).any(0))

np.allclose(locations, locations2)

True

This shows that your output is exactly the same as this output, without the need of explicitly mentioning b[0], b[1]... or using a for loop.

Explanation -

Broadcasting an operation can help you in this case. What you are trying to do is to compare each of the (3,4) matrix elements to each value in b which is (3,). This means that the resultant boolean matrix that you want is going to be three, (3,4) matrices, or (3,3,4)
Once you have done that, you want to take an ANY or OR between the three (3,4) matrices element-wise. That would reduce the (3,3,4) to a (3,4)
Finally you want to use np.nonzero to identify the locations where values are equal to TRUE

The above 3 steps can be done as follows -

Broadcasting comparison operation:

a[None,:,:]==b[:,None,None]] #(1,3,4) == (3,1,1) -> (3,3,4)

Reduction using OR logic:

(a[None,:,:]==b[:,None,None]).any(0)  #(3,3,4) -> (3,4)

Get non-zero locations:

np.nonzero((a[None,:,:]==b[:,None,None]).any(0))

edited Feb 21, 2021 at 21:39

answered Feb 21, 2021 at 20:24

Akshay Sehgal

19.4k3 gold badges26 silver badges57 bronze badges

3 Comments

Alex Shroyer Over a year ago

This does not tell me where these items appear in a.

Akshay Sehgal Over a year ago

You can completely vectorize the above operation. I am adding an update to my answer

Alex Shroyer Over a year ago

Thanks for the update and the explanation!

Alex Shroyer · Accepted Answer · 2021-02-23 14:19:59Z

0

numpy.isin works on multi-dimensional a and b.

In [1]: import numpy as np

In [2]: a = np.random.randint(0, 16, size=(3, 4)); a
Out[2]:
array([[12,  2, 15, 11],
       [12, 15,  5, 10],
       [ 4,  2, 14,  7]])

In [3]: b = [2, 4, 5, 12]

In [4]: c = [[2, 4], [5, 12]]

In [5]: np.isin(a, b).astype(int)
Out[5]:
array([[1, 1, 0, 0],
       [1, 0, 1, 0],
       [1, 1, 0, 0]])

In [6]: np.isin(a, c).astype(int)
Out[6]:
array([[1, 1, 0, 0],
       [1, 0, 1, 0],
       [1, 1, 0, 0]])

In [7]: a[np.isin(a, b)] = 99; a
Out[7]:
array([[99, 99, 15, 11],
       [99, 15, 99, 10],
       [99, 99, 14,  7]])

answered Feb 23, 2021 at 14:19

Alex Shroyer

3,8592 gold badges33 silver badges57 bronze badges

Collectives™ on Stack Overflow

Is there an array method for testing multiple equality values?

3 Answers 3

1 Comment

EDIT:

Explanation -

3 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

EDIT:

Explanation -

3 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related