CODEWITHSUNDEEP
pythonarraysnumpy
Nicolas Gervais
Nicolas Gervais

Reputation: 36704

For all elements in a Numpy array, return their index in another array

Given an array:

big_array = np.array(['dog', 'cat', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog'])

I want to get the position of each of these elements, in a smaller array:

small_array = np.array(['dog', 'cat'])

It should return:

[0, 1, 0, 0, 0, 1, 1, 0]

It would be the equivalent of:

[np.where(i == small_array)[0][0] for i in big_array]

Can it be done without list comprehension, preferably with a Numpy function?

Upvotes: 0

Views: 103

Answers (3)

wwii
wwii

Reputation: 23783

Compare with broadcasting and find index of True in the result along the last axis.

>>> a = np.array(['dog', 'cat', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog'])
>>> b = np.array(['dog','cat'])
>>> c = a[:,None] == b
>>> c.argmax(axis=-1)
array([0, 1, 0, 0, 0, 1, 1, 0], dtype=int64)

>>> a[:,None] == b
array([[ True, False],
       [False,  True],
       [ True, False],
       [ True, False],
       [ True, False],
       [False,  True],
       [False,  True],
       [ True, False]])

Upvotes: 0

Austin
Austin

Reputation: 26037

There's is no numpy function as far as I know, but you could also do a combination of argsort and searchsorted something like:

import numpy as np

big_array = np.array(['dog', 'cat', 'dog', 'dog', 'dog', 'cat', 'cat', 'dog'])

small_array = np.array(['dog', 'cat'])

sortd = np.argsort(small_array)
res = sortd[np.searchsorted(small_array[sortd], big_array)]

print(res)
# [0 1 0 0 0 1 1 0]

Upvotes: 2

John Zwinck
John Zwinck

Reputation: 249572

If you're willing to sort small_array, this will do it:

small_array.sort() # in-place, or `x = np.sort(small_array)` for a sorted copy
np.searchsorted(small_array, big_array)

Upvotes: 2

Related Questions