python - NumPy 通过索引获取多行的快速方法

Question

给定一个 shape 的多维数组(6, 100, 2)和一个索引列表（例如[1,2,3,1,5,0,0,0,1]：），我想获得一个 numpy ndarray，其中包含这些索引中的数据（意思是，输出形状是(9, 100, 2)，行 6,7,8 是彼此）。

天真的解决方案：

arr = []
indexes = [1,2,3,1,5,0,0,0,1]
for i in indexes:
  arr.append(data[i])
arr = np.stack(arr)

这种幼稚的解决方案浪费了大量的时间__getitem__和stack操作，我想知道是否有更快的方法来使用 numpy 做到这一点？

标签： pythonnumpy

Answer 1

你可以只使用下标。因此，如果data是您的 6×100×2 矩阵，并且indexes是您的索引列表，您可以使用以下命令检索结果：

result = data[indexes]