NumPy 多维数组索引答案

【问题标题】：NumPy multi-dimensional array indexingNumPy 多维数组索引
【发布时间】：2018-04-05 18:02:46
【问题描述】：

假设以下示例：

>>> 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 中的哪个元素。

现在我想通过索引idx 从第一个轴 (K) 中选择一个形状为 (J=10, I=200) 的数组。

我怎样才能做到这一点？

【问题讨论】：

标签： python-3.x numpy indexing

【解决方案1】：

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

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

跳过结尾的 : 给我们一个更短的版本 -

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

【讨论】：