NumPy多维数组索引

Question

假设以下示例：

>>> 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)的数组。

我该如何做到这一点？

Answer 1

我们使用idx沿第一个轴索引，同时沿第二个轴选择每个元素，沿着最后一个轴选择。因此，我们可以使用advanced-indexing，就像这样 -

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

跳过尾随:会给我们一个较短的版本 -

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