Replace value in an array by its index in a list

Question

In a given array I want to replace the values by the index of this value in an other array (which doesn't contain duplicates). Here is a simple example of I'm trying to do:

import numpy as np
from copy import deepcopy

a = np.array([[0, 1, 2], [2, 1, 3], [0, 1, 3]])
chg = np.array([3, 0, 2, 1])

b = deepcopy(a)
for new, old  in enumerate(chg):
   b[a == old] = new

print b
# [[1 3 2] [2 3 0] [1 3 0]]

But I need to do that on large arrays so having an explicit loop is not acceptable in terms of execution time.

I can't figure out how to do that using fancy numpy functions...

Answer 1

take is your friend.

a = np.array([[0, 1, 2], [2, 1, 3], [0, 1, 3]])
chg = np.array([3, 0, 2, 1])
inverse_chg=chg.take(chg)
print(inverse_chg.take(a))

gives :

[[1 3 2]
 [2 3 0]
 [1 3 0]]

or more directly with fancy indexing: chg[chg][a], but inverse_chg.take(a) is three times faster.

Answer 2

You can convert chg to a 3D array by adding two new axes at the end of it and then perform the matching comparison with a, which would bring in NumPy's broadcasting to give us a 3D mask. Next up, get the argmax on the mask along the first axis to simulate "b[a == old] = new". Finally, replace the ones that had no matches along that axis with the corresponding values in a. The implementation would look something like this -

mask = a == chg[:,None,None]
out = mask.argmax(0)
invalid_pos = ~mask.max(0)
out[invalid_pos] = a[invalid_pos]