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.
CodePudding user response:
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])
CodePudding user response:
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]