python - 有没有更快的方法在 numpy 中进行这种索引？

Question

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

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

需要一个形状为 (100,) 的 c，这样，

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 毫秒(根据 IPython 的 %timeit)，即每次索引 31 ns。