我有如下所示的 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 循环。它可以工作,但仍然很昂贵。有没有更有效的方法?
最佳答案
是的,你可以这样做:
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。
关于python - 有没有更快的方法在 numpy 中进行这种索引?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/24640689/