Setting value of numpy array based on multiple criteria

Question

I am trying to set the values in a numpy array to zero if it is equivalent to any number in a list.

Lets consider the following array

  a = numpy.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])

I want to set multiple elements of a which are in the list [1, 2, 8] to 0.

The result should be

   [[0, 0, 3],
    [4, 0, 6],
    [7, 0, 9]]

For a single element it's simple

   a[a == 1] = 0

The above only works for a single integer. How it could work for a list?

Answer 1

Using np.in1d you could do the following:

>>> a = np.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])
>>> np.in1d(a, [1, 2, 8])
array([ True,  True, False, False,  True, False, False,  True, False], dtype=bool)
>>> a[np.in1d(a, [1, 2, 8]).reshape(a.shape)] = 0
>>> a
array([[0, 0, 3],
       [4, 0, 6],
       [7, 0, 9]])

Answer 2

Combining my comment to the original question regarding np.where and the excellent answer by @Jaime above using np.in1d:

import numpy as np
a = np.array([[1, 2, 3], [4, 8, 6], [7, 8, 9]])
a = np.where(np.in1d(a, [1,2,8]).reshape(a.shape), 0, a)

EDIT Looks like Jaime's solution is slightly faster:

In [3]: %timeit a[np.in1d(a, [1, 2, 8]).reshape(a.shape)] = 0
10000 loops, best of 3: 45.8 µs per loop
In [4]: %timeit np.where(np.in1d(a, [1,2,8]).reshape(a.shape), 0, a)
10000 loops, best of 3: 66.7 µs per loop