Fastest way to replace all elements in a numpy array

Question

I am trying to find the fastest way to replace elements from a numpy array with elements from another one using a mapping rule. I will give you an example because I think it will be most clearly in this way.

Let's say we have these 3 arrays:

data = np.array([[1,2,3], [4,5,6], [7,8,9], [1,2,3], [7,5,6]])
map = np.array([0, 0, 1, 0, 1])
trans = np.array([[10,10,10], [20,20,20]])

The map array specifies the desired correspondence between data and trans .

The result I want to get is:

array([[11, 12, 13], [14, 15, 16], [27, 28, 29], [11, 12, 13], [27, 25, 26]])

Each element in the array written above is the sum between the element in data and the corresponding element in trans .

I am trying to avoid for loops , because in reality my data and trans arrays are way larger, but couldn't figure out the appropriate vectorised function.

Would you please give me some help?

Answer 1

Index into trans with the indices from map to select rows off trans based on the indices, that act as the row indices for selection and then simply add with data -

data+trans[map]

To do in-situ edit -

data += trans[map]

Word of caution : I would use another variable name than map for the mapping array, which is also a Python built-in to avoid any undesired behavior.

Sample run -

In [23]: data = np.array([[1,2,3], [4,5,6], [7,8,9], [1,2,3], [7,5,6]])
    ...: map1 = np.array([0, 0, 1, 0, 1])
    ...: trans = np.array([[10,10,10], [20,20,20]])
    ...: 

In [24]: trans[map1]
Out[24]: 
array([[10, 10, 10],
       [10, 10, 10],
       [20, 20, 20],
       [10, 10, 10],
       [20, 20, 20]])

In [25]: data + trans[map1]
Out[25]: 
array([[11, 12, 13],
       [14, 15, 16],
       [27, 28, 29],
       [11, 12, 13],
       [27, 25, 26]])