NumPy：按选定的第二轴元素对ndarray第一轴进行排序答案

【问题标题】：NumPy: sort ndarray 1st axis by selected 2nd axis elementNumPy：按选定的第二轴元素对ndarray第一轴进行排序
【发布时间】：2018-06-14 03:05:47
【问题描述】：

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

import numpy as np 
a = np.random.rand(27).reshape(3,3,3)
print(a)

for i in np.arange(np.shape(a)[0]):
    idx = np.argsort(a[i,:,1])
    a[i]=a[i,idx]

print(a)

【问题讨论】：

发布的解决方案对您有用吗？
是的。这正是我想要的。谢谢！

标签： python performance numpy optimization

【解决方案1】：

方法#1

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

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

【讨论】：