Efficiently index a multidemnsional numpy array by another array

Question

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)

Or x[idxs[:,0],idxs[:,1]]? Look here for more info - docs.scipy.org/doc/numpy/reference/… — Divakar
– Divakar, Commented Jul 22, 2016 at 22:05
@Divakar Haha, everytime I post the page refreshes and you've just written it :) — Alex
– Alex, Commented Jul 22, 2016 at 22:09

Alex · Accepted Answer · 2016-07-22 22:11:48Z

3

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

edited Jul 22, 2016 at 22:11

answered Jul 22, 2016 at 22:07

Alex

19.2k9 gold badges65 silver badges82 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Efficiently index a multidemnsional numpy array by another array

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related