Numpy: Indices of multiple values

Question

I would like to set some values of a 2D array to a specific number by indexing them efficiently.

Say I have a 2D numpy array,

A = array([[1, 6, 6],
           [9, 7, 7], 
           [10, 2, 2]])

and I would like to get the indices in the array that belong to a set of numbers, say indList=[10, 1] so that I can set them to zero. However, indList can be a huge list.

Is there a faster way for doing this without a for loop?

As a for loop it would be,

indList = [10, 1]
for i in indList:
    A[A==i] = 0

But this can get inefficient when indList is large.

Answer 1

With numpy, you can vectorize this by first finding the indices of elements that are in indList and then setting them to be zero.

A = np.array([[1, 6, 6],
              [9, 7, 7],
              [10 ,2 ,2]])

A[np.where(np.isin(A, [10,1]))] = 0

This gives

A = [[0 6 6]
     [9 7 7]
     [0 2 2]]

Answer 2

From @Miket25's answer, there is actually no need to add the np.where layer. np.isin(A, [10, 1]) returns a boolean array which is perfectly acceptable as an index. So simply do

A[np.isin(A, [10, 1])] = 0