Fast vectorized indexing in numpy

Question

Suppose we have an array of indices of another numpy array:

import numpy as np
a = np.array([0, 3, 1])
b = np.array([0, 10, 20, 30, 40, 50, 60, 70])

We can use the array a as index directly:

b[a] # np.array([0, 30, 10])

But what if array a has more than one dimension? For example,

a = np.array([[0, 2], [1, 3], [2, 4]])
# I want to get b[a] = np.array([[0, 20], [10, 30], [20, 40]])

Numpy indexing doesn't work if the number of dimensions of a is greater than 1. We can achieve the desired result by using map

map(lambda x: b[x], a)

However, it is quite slow. For 1-dimensional case, direct indexing is about 10-100 times faster than using map.

Is there a way to do it faster?

Answer 1

What's the problem? I can index b with a 2d array. The output just matches a1 in shape:

In [64]: b = np.array([0, 10, 20, 30, 40, 50, 60, 70])
In [65]: a1 = np.array([[0, 2], [1, 3], [2, 4]])
In [66]: b[a1]
Out[66]: 
array([[ 0, 20],
       [10, 30],
       [20, 40]])

b[a1] is not the same as b[a1[:,0],a1[:,1]]. That is the 2 columns of a1 do not provide two indices (which would require a 2d b).