You can use np.in1d
-
A*np.in1d(A,[1,3,4]).reshape(A.shape)
Also, np.where
could be used -
np.where(np.in1d(A,[1,3,4]).reshape(A.shape),A,0)
You can also use np.searchsorted
to find such matches by using its optional 'side'
argument with inputs as left
and right
and noting that for the matches, the searchsorted would output different results with these two inputs. Thus, an equivalent of np.in1d(A,[1,3,4])
would be -
M = np.searchsorted([1,3,4],A.ravel(),'left') !=
np.searchsorted([1,3,4],A.ravel(),'right')
Thus, the final output would be -
out = A*M.reshape(A.shape)
Please note that if the input search list is not sorted, you need to use the optional argumentsorter
with its argsort
indices in np.searchsorted
.
Sample run -
In [321]: A
Out[321]:
array([[1, 1, 0, 2, 2],
[1, 1, 0, 2, 0],
[0, 0, 0, 0, 0],
[3, 3, 0, 4, 4],
[3, 3, 0, 4, 4]])
In [322]: A*np.in1d(A,[1,3,4]).reshape(A.shape)
Out[322]:
array([[1, 1, 0, 0, 0],
[1, 1, 0, 0, 0],
[0, 0, 0, 0, 0],
[3, 3, 0, 4, 4],
[3, 3, 0, 4, 4]])
In [323]: np.where(np.in1d(A,[1,3,4]).reshape(A.shape),A,0)
Out[323]:
array([[1, 1, 0, 0, 0],
[1, 1, 0, 0, 0],
[0, 0, 0, 0, 0],
[3, 3, 0, 4, 4],
[3, 3, 0, 4, 4]])
In [324]: M = np.searchsorted([1,3,4],A.ravel(),'left') !=
...: np.searchsorted([1,3,4],A.ravel(),'right')
...: A*M.reshape(A.shape)
...:
Out[324]:
array([[1, 1, 0, 0, 0],
[1, 1, 0, 0, 0],
[0, 0, 0, 0, 0],
[3, 3, 0, 4, 4],
[3, 3, 0, 4, 4]])
Runtime tests and verify outputs -
In [309]: # Inputs
...: A = np.random.randint(0,1000,(400,500))
...: lst = np.sort(np.random.randint(0,1000,(100))).tolist()
...:
...: def func1(A,lst):
...: return A*np.in1d(A,lst).reshape(A.shape)
...:
...: def func2(A,lst):
...: return np.where(np.in1d(A,lst).reshape(A.shape),A,0)
...:
...: def func3(A,lst):
...: mask = np.searchsorted(lst,A.ravel(),'left') !=
...: np.searchsorted(lst,A.ravel(),'right')
...: return A*mask.reshape(A.shape)
...:
In [310]: np.allclose(func1(A,lst),func2(A,lst))
Out[310]: True
In [311]: np.allclose(func1(A,lst),func3(A,lst))
Out[311]: True
In [312]: %timeit func1(A,lst)
10 loops, best of 3: 30.9 ms per loop
In [313]: %timeit func2(A,lst)
10 loops, best of 3: 30.9 ms per loop
In [314]: %timeit func3(A,lst)
10 loops, best of 3: 28.6 ms per loop