使用Numpy索引三维阵列

Question

我有一个三维(x，y，z(数组和一个索引向量。该向量的大小等于数组的维度x。它的目标是索引一个特定的y，使其相应的z，即预期结果具有维度(x，z(。

我写了一段按预期工作的代码，但有人知道Numpy函数是否可以取代for循环并更优化地解决问题吗？

arr = np.random.rand(100,5,2)
result = np.random.rand(100,2)
id = [np.random.randint(0, 5) for _ in range(100)]
for i in range(100): 
result[i] = arr[i,id[i]]

Answer 1

您可以通过以下代码实现这一点：

import numpy as np
arr = np.random.randn(100, 5, 2)
ids = np.random.randint(0, 5, size=100)
res = arr[range(100), ids]
res.shape # (100, 2)