Numpy高级索引

Question

我有以下形状的数组

voxel_grids : (128, 32, 32, 32)
indices : (128, 3, 1024)

我想构建一个数组scalar (128, 1024)，以便

scalar[i,j] = voxel_grids[i, indices[i, 0, j], indices[i, 1, j], indices[i,2,j]]

使用numpy（高级）索引是否有直接的方法来执行此操作？

Answer 1

你可以这样做 -

m = voxel_grids.shape[0]
out = voxel_grids[np.arange(m)[:,None],indices[:,0],indices[:,1],indices[:,2]]

其他方法是将indices中的三个切片提取为三个变量并将其用于索引。这可能不比前一个更有效，但可能会更容易一些。它显示如下 -

m = voxel_grids.shape[0]
x,y,z = indices.swapaxes(0,1)
out = voxel_grids[np.arange(m)[:,None],x,y,z]