索引一个numpy数组[重复]

Question

考虑以下场景。

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]

使用 numpy 索引实现此行为的推荐（就速度而言）方法是什么？

Answer 1

这正是np.choose的目的，你只需要先重新排列y轴。

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

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

它比循环方法快 15 倍。