Mike Mike - 4 months ago 17
Python Question

Efficiently index a multidemnsional numpy array by another array

I have an array

x
which specific values I would like to access, whose indices are given by another array.

For example,
x
is

array([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]])


and the indices are an array of Nx2

idxs = np.array([[1,2], [4,3], [3,3]])


I would like a function that returns an array of x[1,2], x[4,3], x[3,3] or [7, 23, 18]. The following code does the trick, but I would like to speed it up for large arrays, perhaps by avoiding the for loop.

import numpy as np

def arrayvalsofinterest(x, idx):
output = np.zeros(idx.shape[0])
for i in range(len(output)):
output[i] = x[tuple(idx[i,:])]
return output

if __name__ == "__main__":
xx = np.arange(25).reshape(5,5)
idxs = np.array([[1,2],[4,3], [3,3]])
print arrayvalsofinterest(xx, idxs)

Answer

You can pass in an iterable of axis0 coordinates and an iterable of axis1 coordinates. See the Numpy docs here.

i0, i1 = zip(*idxs)
x[i0, i1]

As @Divakar points out in the comments, this is less memory efficient than using a view of the array i.e.

x[idxs[:, 0], idxs[:, 1]]