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.
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])
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]
np.where(condition) is actually a shortcut for np.asarray(condition).nonzero() ;)
