python - 使用 numpy 进行多维索引的最佳方法是什么？

Question

我正在尝试对 3D numpy 数组进行一些索引。基本上我有一个数组 phi ，其形状为 (F,A,D);例如(5,3,7)。生成后例如如下:

F=5; A=3; D=7; phi = np.random.random((F,A,D))

我的目标是能够使用二维数组(例如[[0,1,2],[5)对A和D进行索引,5,6]]，表示A中的第一个位置取第三维中以0索引的值，A中的第一个位置取第三维中以1索引的值A 的第二个位置，依此类推。结果的形状应为 (F,A,2) 或 (F,2,A)。

这相当于手动循环“索引器数组”的所有值，例如:

phi[:,0,0]; phi[:,1,1]; phi[:,2,2]
phi[:,0,5]; phi[:,1,5]; phi[:,2,6]

直觉上我会做类似 phi[:,:,[[0,1,2],[3,3,3]]] 的事情，但它的形状最终是 (5,3,2,3)。

关于如何获得正确结果有什么想法吗？

Answer 1

我想这就是你想要的

phi[:,range(A),[[0,1,2],[5,5,6]]]

你的尝试

phi[:,:,[[0,1,2],[5,5,6]]]

对前两个维度的每个值沿第三个维度获取值，因此最终得到的形状为(5,3,2,3)。

但是，根据您的示例，您希望第二个维度不断增加，这是通过 range(A) 和 numpy 的广播在我的代码中完成的。