Find the index position of all elements within an array with matching values from a test array in numpy

Question

I have a numpy array like:

a = np.array([[1, 2, 3, 456], [2, 3, 4, 789], [3, 4, 5, 101112], [4, 5, 6, 131415]])

I have an array of numbers like:

b = np.array([101112, 456])

I am looking for:

[2, 0]

How can I get the index positions in a using b?

Currently, I am using a nested loop which is highly inefficient.

I cannot get np.where to do this, at least with my limited understanding.

Answer 1

Given your arrays you can ask:

a = np.array([[1, 2, 3, 456], [2, 3, 4, 789], [3, 4, 5, 101112], [4, 5, 6, 131415]])
b = np.array([101112, 456])

np.isin(a, b)

and get:

array([[False, False, False,  True],
       [False, False, False, False],
       [False, False, False,  True],
       [False, False, False, False]])

Passing that to any() and then to argwhere will give you the one-liner:

np.argwhere(np.any(np.isin(a, b), axis=1)).ravel()
# array([0, 2])

Answer 2

You can apply np.where for each value of b

import numpy as np

a = np.array([[1, 2, 3, 456], [2, 3, 4, 789], 
              [3, 4, 5, 101112], [4, 5, 6, 131415]])
b = np.array([101112, 456])
print([np.where(a == v)[0][0] for v in b])  # [2, 0]