Python 3.x NumPy多维数组索引

假设以下示例：

>>> a = np.random.randint(0, 10, (3, 10, 200))
>>> print(a.shape)
(3, 10, 200)
>>> 
>>> idx = np.random.randint(0, 3, 10)
>>> print(idx)
[2, 0, 0, 0, 1, 2, 1, 2, 0, 0]

a

是一个形状数组

（K=3，J=10，I=200）

idx

是一个与

a.shape[1]

长度相同的数组，即包含J=10个元素。每个索引表示应选择K的哪个元素

现在，我想通过索引从第一个轴（K）中选择

idx

，以获得形状数组

（J=10，I=200）

---

如何实现这一点？

我们使用

idx

沿第一个轴进行索引，同时沿第二个轴和最后一个轴选择每个元素。因此，我们可以这样使用-

a[idx, np.arange(len(idx)),:]

跳过尾随的

：

可以得到一个较短的版本-

a[idx, np.arange(len(idx))]