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]