这是我在这里回答的动机。
给定形状为（n0，n1）的数组A和形状为（n0）的数组J，我想创建形状为（n0）的数组B
B[i] = A[i,J[i]]
我也想将其推广到k维数组，其中A具有形状(n0,n1,...,nk)，J具有形状(n0,n1,...,n(k-1))
有一些混乱的，扁平化的方法可以对索引顺序进行假设：
import numpy as np
B = A.ravel()[ J+A.shape[-1]*np.arange(0,np.prod(J.shape)).reshape(J.shape) ]
问题是，有没有一种方法不依赖于展平数组和手动处理索引？
对于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 2arange索引的广播尺寸与相同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 loop
In [614]: timeit Abig[np.ix_(*[np.arange(x) for x in Jbig.shape])+(Jbig,)]
10000 loops, best of 3: 22.7 µs per loop
In [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。
本文收集自互联网，转载请注明来源。
如有侵权，请联系 [email protected] 删除。
我来说两句