Mike - 1 year ago 92

Python Question

I have an array

`x`

For example,

`x`

`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 Source

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