CODEWITHSUNDEEP
pythonarraysnumpy
Mike
Mike

Reputation: 1827

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)

Upvotes: 0

Views: 93

Answers (1)

Alex
Alex

Reputation: 19104

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]]

Upvotes: 3

Related Questions