使用 Numpy 索引 3D 数组

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)