Indexing a numpy array

Question

Consider the following scenario.

import numpy as np
x = np.random.randint(0,21,size=(10,64,64))
y = np.random.rand(10,21,64,64)

z = np.empty((10,64,64))

for i in range(10):
    for j in range(64):
        for k in range(64):
            z[i][j][k] = y[i][x[i][j][k]][j][k]

What is the recommended (in terms of speed) way to implement this behavior using numpy indexing?

https://stackoverflow.com/questions/46103044/index-n-dimensional-array-with-n-1-d-array — user2357112, Dec 05 '17 at 02:03

score 1 · Accepted Answer · answered Dec 05 '17 at 07:22

1

It is exactly the aim of np.choose, You just need to rearange y axis first.

In [6]: z2=np.choose(x,np.rollaxis(y,1))

In [7]: np.allclose(z,z2)
Out[7]: True

It's 15x faster than the loop method.

answered Dec 05 '17 at 07:22

B. M.

18,243
2
35
54

Indexing a numpy array

1 Answers1