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.

Asked By: Alex Shroyer

||

Source

Answer 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))

Answered By: Akshay Sehgal

Answer 2

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 By: ALollz

Answer 3

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 By: Alex Shroyer

Is there an array method for testing multiple equality values?

Question:

Answers:

EDIT:

Explanation –