find indices of matching rows in 2d array without loop (one shot)

user_13829734 0
Viewed 753

I want to find multiple matching rows from 2d array a

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

I have to search following entries

 b = np.array([[4,6],
          [4,7]])

I know that I can loop over b and do following

for i in range(len(b)) :
   print(np.where(np.all(a==b[i],axis=1))[0])

And I get following

[2]
[4]

Can I get [[2],[4]] directly without using any loop?

2d numpy search arrays python
2 Answers

If you want the indices you will generally reach for the arg_x functions like argmax and argwhere. Here np.argwhere will give you the indices if you can figure out how to pass the correct list of booleans. You can do that with np.isin():

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

b = np.array([[4,6], [4,7]])

np.argwhere(np.isin(a, b).all(axis=1))

Which returns:

array([[2],
       [4]])
user_38746230

This should be a fast solution noticing that the two pairs have the same first coordinate:

np.where((a[:, 0] == 4) & ((a[:, 1] == 6) | (a[:, 1] == 7)))  
# Out:
# (array([2, 4]),)

The expression

print((a[:, 0] == 4) & ((a[:, 1] == 6) | (a[:, 1] == 7)))

gives

[False False  True False  True False]
user_23147370
Related