1

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.

matttree
  • 125
  • 1
  • 1
  • 11

1 Answers1

2

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])
hpaulj
  • 221,503
  • 14
  • 230
  • 353
  • That first solution is truely elegant - I have Never seen seen the ```None``` index comparison before, and will definitely use in the future. Can you place a reference for how it works? – D A Dec 09 '21 at 18:18