在多维 numpy 数组中索引多个元素

Question

我想使用另一个索引数组提取给定多维 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]

？

Answer 1

您正在寻找 numpy 高级索引

https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing

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

idx

：

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