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

idx is an array of the same length as a.shape[1], i.e. contains J = 10 elements. Each index denotes which element of K should be chosen.

现在，我想通过索引idx从第一个轴(K)中进行选择，以得到形状为(J=10, I=200)的阵列.

Now I'd like to select from the first axis (K) by the indices idx to get an array of shape (J=10, I=200) back.

我该怎么做?

推荐答案

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

We are using idx to index along the first axis, while selecting per element along the second axis and all along the last one. Thus, we can use advanced-indexing, like so -

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

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

Skipping the trailing : gives us a shorter version -

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

这篇关于NumPy多维数组索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！