Question

我有一个数组n x m x k。任务是通过选定的第二轴元素对第一轴进行排序。目前，我在下面的例子中这样做。输入数组为3x3x3，轴1按第2轴的第1个元素排序。是否有更有效的方式使用nditer？

/Users/test/jenkinsslave/tools/hudson.plugins.sonar.SonarRunnerInstallation/SonarQube_Scanner_2.8/bin/sonar-scanner -X -e -Dsonar.host.url=http://sonarqube:9090 -Dsonar.projectBaseDir=./../module2

Answer 1

方法＃1

获取argsort索引一次性选择沿第一轴的所有元素，然后使用NumPy的advanced-indexing沿第二轴索引以获得重新排列的输出 - < / p>

idx = np.argsort(a[:,:,1])
a_out = a[np.arange(a.shape[0])[:,None], idx]

方法＃2

如果我们在输入数组中进行原位编辑（将结果写回输入），我们可以在预先计算那些argsort索引后运行循环，如此 -

idx = np.argsort(a[:,:,1])
for i,indx in enumerate(idx):
    a[i] = a[i,indx]

基准

# Original method
In [130]: np.random.seed(0)
     ...: a = np.random.rand(100,100,100)

In [131]: %%timeit
     ...: for i in np.arange(np.shape(a)[0]):
     ...:     idx = np.argsort(a[i,:,1])
     ...:     a[i]=a[i,idx]
1000 loops, best of 3: 1.63 ms per loop

# Approach #1
In [132]: np.random.seed(0)
     ...: a = np.random.rand(100,100,100)

In [133]: %%timeit
     ...: idx = np.argsort(a[:,:,1])
     ...: a_out = a[np.arange(a.shape[0])[:,None], idx]
1000 loops, best of 3: 1.6 ms per loop

# Approach #2
In [134]: np.random.seed(0)
     ...: a = np.random.rand(100,100,100)

In [135]: %%timeit
     ...: idx = np.argsort(a[:,:,1])
     ...: for i,indx in enumerate(idx):
     ...:     a[i] = a[i,indx]
1000 loops, best of 3: 1.24 ms per loop

NumPy：通过选择的第二轴元素对ndarray第一轴进行排序

1 个答案:

基准