是否有更快的方法在numpy中进行这种索引？

Question

我有如下所示的numpy数组a和b。

a.shape = (100,10000)
b.shape = (100,)

需要c形状（100，），

for i in range(0,len(b)):
   c[i] = a[i,b[i]]

我希望以更快的方式完成上述操作。 For循环不能在我正在使用的场景中使用..它可以工作，但仍然很昂贵。有没有更有效的方法呢？

Answer 1

是的，你可以这样做：

import numpy as np

c = a[np.arange(len(b)), b]

这和它一样快。在以下代码中

import numpy as np

a = np.random.random((1000000, 10))
b = np.random.randint(0,10,1000000)
c = a[np.arange(len(b)), b]

最后一次操作需要31 ms（通过IPython的%timeit），即每次索引为31 ns。