numpy where implementation for this case

user_9790945 0
Viewed 54

I have an array a=np.arange(0,10,1). I am given with another array say b=[1,3,0]. I need to find all the indices corresponding to location of all elements of b in the original array a. The answer in this case should be [1,3,0] itself. How would I do this using np.where. I know for one element I could do something like np.where(a==1), but I want to know if there is an efficient way for an entire array of comparisons.

numpy python
1 Answers

A straight forward approach using broadcasted equality:

In [104]: a=np.arange(10)
In [105]: a[:,None]==[1,3,0]
Out[105]: 
array([[False, False,  True],
       [ True, False, False],
       [False, False, False],
       [False,  True, False],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False]])
In [106]: _.any(axis=1)
Out[106]: 
array([ True,  True, False,  True, False, False, False, False, False,
       False])
In [107]: np.nonzero(_)
Out[107]: (array([0, 1, 3]),)

Another way to get the boolean array (a little slower for this sample):

In [110]: np.in1d(a,[1,3,10])
Out[110]: 
array([False,  True, False,  True, False, False, False, False, False,
       False])
user_9019250
Related