element-wise lookup on one ndarray to another ndarray of different shapes

Question

Hi,

I am new to numpy. Am wonder is there a way to do lookup of two ndarray of different shapes? for example, i have 2 ndarrays as below:

X = array([[0, 3, 6],
       [3, 3, 3],
       [6, 0, 3]])

Y = array([[0, 100],
       [3, 500],
       [6, 800]])

and would like to lookup each element of X in Y, then be able to return the second column of Y:

Z = array([[100, 500, 800],
       [500, 500, 500],
       [800, 100, 500]]) 

thanks, fahhean

Answer 1

You can directly use NumPy's efficient array operations:

Y_dict = dict(Y)
Z = vectorize(lambda x: Y_dict[x])(X)

This approach has the advantage of working whatever the dimension of X (1-dimensional array, 2- or N-dimensional array…).

The vectorized function automatically applies the dictionary look-up to each element of array X in turn.

The first line is just there for optimization purposes; otherwise the dictionary construction would take place for every element in X.