numpy.take由于某种原因,它比花式索引要快得多。唯一的技巧是将数组视为平面。
In [1]: a = np.random.randn(12,6).astype(np.float32)In [2]: c = np.random.randint(0,6,size=(1e5,12)).astype(np.uint8)In [3]: r = np.arange(12)In [4]: %timeit a[r,c].sum(-1)10 loops, best of 3: 46.7 ms per loopIn [5]: rr, cc = np.broadcast_arrays(r,c)In [6]: flat_index = rr*a.shape[1] + ccIn [7]: %timeit a.take(flat_index).sum(-1)100 loops, best of 3: 5.5 ms per loopIn [8]: (a.take(flat_index).sum(-1) == a[r,c].sum(-1)).all()Out[8]: True
我认为,除此以外,您还将看到速度提高很多的唯一其他方法是使用PyCUDA之类的东西为GPU编写自定义内核。



