Find where a NumPy array is equal to any value in a list of values

Question

I have an array of integers and want to find where that array is equal to any value in a list of multiple values.

This can easily be done by treating each value indi

Answer 1

@Bas's answer is the one you're probably looking for. But here's another way to do it, using numpy's vectorize trick:

import numpy as np
S = set([0,2,6,8])

@np.vectorize
def contained(x):
    return x in S

contained(fake)
=> array([[ True, False,  True],
          [False, False, False],
          [ True, False,  True]], dtype=bool)

The con of this solution is that contained() is called for each element (i.e. in python-space), which makes this much slower than a pure-numpy solution.

Answer 2

NumPy 0.13+

As of NumPy v0.13, you can use np.isin, which works on multi-dimensional arrays:

>>> element = 2*np.arange(4).reshape((2, 2))
>>> element
array([[0, 2],
       [4, 6]])
>>> test_elements = [1, 2, 4, 8]
>>> mask = np.isin(element, test_elements)
>>> mask
array([[ False,  True],
       [ True,  False]])

NumPy pre-0.13

The accepted answer with np.in1d works only with 1d arrays and requires reshaping for the desired result. This is good for versions of NumPy before v0.13.

Answer 3

The function numpy.in1d seems to do what you want. The only problems is that it only works on 1d arrays, so you should use it like this:

In [9]: np.in1d(fake, [0,2,6,8]).reshape(fake.shape)
Out[9]: 
array([[ True, False,  True],
       [False, False, False],
       [ True, False,  True]], dtype=bool)

I have no clue why this is limited to 1d arrays only. Looking at its source code, it first seems to flatten the two arrays, after which it does some clever sorting tricks. But nothing would stop it from unflattening the result at the end again, like I had to do by hand here.