NumPy 3D数组索引

Question

我有一个3D numpy数组（n_samples x num_components x 2），在下面的例子中，n_samples=5，num_components=7。

我还有另一个数组（indices），它是每个样本所选组件，其形状为（n_samples,）。

我想根据索引从数据数组中选择数据，以便生成结果数组为n_samples x 2。

以下是代码：

import numpy as np
np.random.seed(77)
data=np.random.randint(low=0, high=10, size=(5, 7, 2))
indices = np.array([0, 1, 6, 4, 5])
#how can I select indices from the data array?

例如，对于数据0，所选的组件应该是第0个，对于数据1，所选的组件应该是第1个。

请注意，我不能使用任何for循环，因为我正在使用Theano，解决方案应仅基于numpy。

- Ash

Answer 1

这是您正在寻找的内容吗？

In [36]: data[np.arange(data.shape[0]),indices,:]
Out[36]: 
array([[7, 4],
       [7, 3],
       [4, 5],
       [8, 2],
       [5, 8]])

- hpaulj

---

是的，就是这样。我认为这是最好的答案。 - Ash

Answer 2

要获取组件#0，请使用：

data[:, 0]

即我们获取轴0（样本）上的每个条目，并且在轴1（组件）上仅获取第0个条目，并隐式地获取其余轴上的所有内容。

这可以很容易地推广到：

data[:, indices]

选择所有相关组件。

---

但 OP 真正想要的只是该数组的对角线，即 (data[0, indices[0]], (data[1, indices[1]]), ...) 高维数组的对角线可以使用 diagonal 函数提取：

>>> np.diagonal(data[:, indices])
array([[7, 7, 4, 8, 5],
       [4, 3, 5, 2, 8]])

您可能需要转置结果。

您可能需要对结果进行转置。

- kennytm

---

data[:,indices][np.arange(n_samples),np.arange(n_samples)] 这段代码可以运行，但是不够清晰易懂。 - jyalim

@Afshin 你觉得 np.diagonal(data[:, indices]).T 怎么样？ - kennytm

我认为这就是我需要的，尽管我认为它有一个简单的索引。如果您更改答案，我将能够选择它作为答案。 - Ash

两个答案都会扩展矩阵并选择对角线部分，我应该测试哪一个更快，以及是否有更好的解决方案，因为在我的情况下，索引是batch_size的数量，可能很大，所以最好不要扩展矩阵。 - Ash

@Afshin 已更新。我猜 hpaulj 的更快，因为它不需要构建中间的 #Samples×#Indices×2 数组。 - kennytm

Answer 3

您有多种方法可以实现这一点，但这是我的循环建议：

selection = np.array([ datum[indices[k]] for k,datum in enumerate(data)])

生成的数组selection具有所需的形状。

- jyalim

---

我不能使用for循环，因为我正在使用Theano，需要一个纯numpy的解决方案。 - Ash

我认为你应该在原始帖子中添加那个约束条件。 - jyalim