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?
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.
Collected from the Internet
Please contact [email protected] to delete if infringement.
Comments