Numpy: Efficient elements of array A in terms of the indices of matching elements of array B

Question

Suppose we have some numpy array A = [1, 2, 2, 3] and its unique array B = [1, 2, 3] where array B is basically the set for array A.

I wish to get a new array C = [0, 1, 1, 2] that is array A but with its elements swapped for the indices of the matching elements in array B.

How would I achieve this?

Note: My arrays A and B contain 18m and 138k values respectively, so the solution must be efficient.

Answer 1

Use:

import numpy as np

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

res = (A[:, None] == B).argmax(1)
print(res)

Output

[0 1 1 2]

Alternative:

_, res = (A[:, None] == B).nonzero()
print(res)

Output

[0 1 1 2]

Alternative, that should use less memory:

import numpy as np

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

# find the uniques values together with the inverse
Bp, inverse = np.unique(A, return_inverse=True)

# find the indices mapping Bp to B
indices = (Bp[:, None] == B).argmax(1)

# index on the mapping
result = indices[inverse]
print(result)

Output

[0 1 1 2]

If B is always the unique sorted elements of A, you can do directly:

import numpy as np

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

# find the uniques values together with the inverse
_, inverse = np.unique(A, return_inverse=True)

print(inverse)

Output

[0 1 1 2]

Answer 2

Using numpy.where :

A = np.array([1, 2, 2, 3])
B = np.array([1, 2, 3])

C = np.where(A==B[:,None])[0]

output: array([0, 1, 1, 2])