How to map a 2D numpy array onto a 1D array with another 2D array as a map?

Question

So I have two 2D numpy arrays, data containing the actual sample data:

[[12 15  5  0]
 [ 3 11  3  7]
 [ 9  3  5  2]
 [ 4  7  6  8]]

and the other, location , containing a map, another 2D array of unique, non-overlapping int values corresponding to the spaces in a new 1D array of the same size as the two:

[[ 5  6  9 10]
 [ 4  7  8 11]
 [ 3  2 13 12]
 [ 0  1 14 15]]

The only way I've been able to run the transfer so far is with a simple for loop:

arr = np.zeros(4*4, dtype = int)

for i in range(4):
    for j in range(4):
        mapval = location[i, j]
        arr[mapval] = data[i, j]

Which does correctly output [ 4 7 3 9 3 12 15 11 3 5 0 7 2 5 6 8]

This is fine with the simple 4*4 array, but the actual dataset clocks in at 512*512, and this method takes rather very long to complete. So my question is, are there any functions, or methods exploiting ufuncs/numpy's fast processing capabilities to get this done more efficiently?

Answer 1

You need the index that sorts the location array to reorder the data array, which can be calculated using argsort :

data.ravel()[location.ravel().argsort()]
# array([ 4,  7,  3,  9,  3, 12, 15, 11,  3,  5,  0,  7,  2,  5,  6,  8])

---

import numpy as np
data = np.array([[12, 15,  5,  0],
 [ 3, 11,  3,  7],
 [ 9,  3,  5,  2],
 [ 4,  7,  6,  8]])

location = np.array([[ 5,  6,  9, 10],
 [ 4,  7,  8, 11],
 [ 3,  2, 13, 12],
 [ 0,  1, 14, 15]])