是否有更快的方法在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毫秒（通過IPython的的%timeit），即每分度31納秒。