如何选择具有（n-1）D个轴索引的nD阵列轴上的值？

如何选择具有（n-1）D个轴索引的nD阵列轴上的值？

对于2d和1d情况，此索引有效：

A[np.arange(J.shape[0]), J]

可以通过重塑为2d（并向后）将其应用于更多尺寸：

A.reshape(-1, A.shape[-1])[np.arange(np.prod(A.shape[:-1])).reshape(J.shape), J]

对于3d，

A

此方法有效：

A[np.arange(J.shape[0])[:,None], np.arange(J.shape[1])[None,:], J]

其中1st 2

arange

索引的广播尺寸与相同

J

。

使用中的函数时

lib.index_tricks

，可以表示为：

A[np.ogrid[0:J.shape[0],0:J.shape[1]]+[J]]A[np.ogrid[slice(J.shape[0]),slice(J.shape[1])]+[J]]

或用于多个维度：

A[np.ix_(*[np.arange(x) for x in J.shape])+(J,)]A[np.ogrid[[slice(k) for k in J.shape]]+[J]]

---

对于小的

A

和

J

（例如，2 * 3 * 4），

J.choose(np.rollaxis(A,-1))

速度更快。所有额外的时间都在准备索引元组中。

np.ix_

比快

np.ogrid

。

np.choose

有大小限制。在最高端它比

ix_

：

In [610]: Abig=np.arange(31*31).reshape(31,31)In [611]: Jbig=np.arange(31)In [612]: Jbig.choose(np.rollaxis(Abig,-1))Out[612]: array([  0,  32,  64,  96, 128, 160, ... 960])In [613]: timeit Jbig.choose(np.rollaxis(Abig,-1))10000 loops, best of 3: 73.1 µs per loopIn [614]: timeit Abig[np.ix_(*[np.arange(x) for x in Jbig.shape])+(Jbig,)]10000 loops, best of 3: 22.7 µs per loopIn [635]: timeit Abig.ravel()[Jbig+Abig.shape[-1]*np.arange(0,np.prod(Jbig.shape)).reshape(Jbig.shape) ]10000 loops, best of 3: 44.8 µs per loop

我在https://stackoverflow.com/a/28007256/901925上做了类似的索引测试，发现

flat

对于更大的数组（例如

n0=1000

），索引速度更快。那是我了解的32个限制的地方

choice

。