我想使用另一个索引数组提取给定多维 numpy 数组的元素。但是它的行为并不符合我的预期。下面是一个简单的例子：

import numpy as np

a = np.random.random((3, 3, 3))
idx = np.asarray([[0, 0, 0], [0, 1, 2]])

b = a[idx]
print(b.shape)  # expect (2, ), got (2, 3, 3, 3)

为什么会这样呢？我应该如何修改代码以使 b 仅包含两个元素：a[0, 0, 0] 和 a[0, 1, 2]？

正确答案

您正在寻找 numpy 高级索引

https://www.php.cn/link/2d661a763280f48803f3c9ba8ba0e00b

在您的情况下，您需要在每个轴上使用 idx：

a[idx[:,0], idx[:, 1], idx[:, 2]].shape == (2,)  # True