我正在使用多维NumPy数组a，它是说2x2矩阵的“向量”。
我想对a进行排序，以便2x2矩阵按行规范排序。
import numpy as np
a = np.array([[[3, 4],
[1, 2]],
[[5, 6],
[7, 8]]])
sortidxs = np.argsort(np.linalg.norm(a, axis=-1))
a = np.array([a[_][sortidxs[_]] for _ in range(a.shape[0])])
# And the final output should be:
print(a)
[[[1 2]
[3 4]]
[[5 6]
[7 8]]]
上面的代码片段完成了我想要的（不完全，请看下面的编辑）。但我一直在寻找避免循环的方法
a = np.array([a[_][sortidxs[_]] for _ in range(a.shape[0])])
- 编辑 -
上面的例子错过了问题的关键部分。a可能有更多的“空”维度，即
a = np.array([[[3, 4],
[1, 2]],
[[5, 6],
[7, 8]]])
a = a.reshape((2,1,2,2))
a现在看起来像：
In [257]: a
Out[257]:
array([[[[3, 4],
[1, 2]]],
[[[5, 6],
[7, 8]]]])
并且在排序之后它应该是
In [259]: a
Out[259]:
array([[[[1, 2],
[3, 4]]],
[[[5, 6],
[7, 8]]]])
a也可以在开头具有以下维度(1,2,2,2)或更多此类“空”维度。我希望这种情况也适用于这些情况。
答案 0 :(得分：3)
您可以使用advanced-indexing -
a[np.arange(a.shape[0])[:,None], sortidxs]
示例运行 -
In [144]: a = np.random.randint(0,9,(2,3,4))
In [145]: a
Out[145]:
array([[[1, 1, 5, 5],
[1, 1, 7, 5],
[6, 1, 2, 8]],
[[7, 2, 5, 4],
[3, 7, 3, 7],
[8, 4, 4, 6]]])
In [146]: sortidxs = np.argsort(np.linalg.norm(a, axis=-1))
In [147]: np.array([a[_][sortidxs[_]] for _ in range(a.shape[0])])
Out[147]:
array([[[1, 1, 5, 5],
[1, 1, 7, 5],
[6, 1, 2, 8]],
[[7, 2, 5, 4],
[3, 7, 3, 7],
[8, 4, 4, 6]]])
In [149]: a[np.arange(a.shape[0])[:,None], sortidxs]
Out[149]:
array([[[1, 1, 5, 5],
[1, 1, 7, 5],
[6, 1, 2, 8]],
[[7, 2, 5, 4],
[3, 7, 3, 7],
[8, 4, 4, 6]]])
进一步提升绩效
我们可以使用np.einsum -
sortidxs
sortidxs = np.einsum('ijk,ijk->ij',a,a).argsort()
让我们花时间验证这个想法 -
In [94]: a = np.random.randint(0,9,(20,30,40))
In [95]: %timeit np.argsort(np.linalg.norm(a, axis=-1))
10000 loops, best of 3: 63.5 µs per loop
In [96]: %timeit np.einsum('ijk,ijk->ij',a,a).argsort()
10000 loops, best of 3: 19.7 µs per loop
In [97]: a = np.random.randint(0,9,(200,300,400))
In [98]: %timeit np.argsort(np.linalg.norm(a, axis=-1))
10 loops, best of 3: 88.6 ms per loop
In [99]: %timeit np.einsum('ijk,ijk->ij',a,a).argsort()
10 loops, best of 3: 22.6 ms per loop
尺寸更高的数组
对于a为4D数组的其他情况，我们需要使用更多数组进行索引。
1]对于第一轴：在最后使用np.arange(a.shape[0])和两个新轴。
2]对于第二轴：使用np.arange(a.shape[0])，最后一个新轴。
3]对于第三轴：使用sortidxs进行索引。
因此，我们会：
m,n,r,s = a.shape
out = a[np.arange(m)[:,None,None],np.arange(n)[:,None], sortidxs]
单例暗淡的数组（dim，长度= 1）
作为一个特例，让我们说输入数组的第二个轴已经是一个单独的轴，我们可以简单地使用0作为那个轴，从而简化了事情，就像这样 -
a[np.arange(m)[:,None,None],0, sortidxs]
示例运行 -
In [58]: a = np.array([[[3, 4],
...: [1, 2]],
...:
...: [[5, 6],
...: [7, 8]]])
...:
...: a = a.reshape((2,1,2,2))
...:
In [59]: sortidxs = np.argsort(np.linalg.norm(a, axis=-1))
In [60]: a[np.arange(a.shape[0])[:,None,None],0, sortidxs]
Out[60]:
array([[[[1, 2],
[3, 4]]],
[[[5, 6],
[7, 8]]]])
为具有通用形状(2,3,4)的数组运行另一个示例以使事情变得非常清晰 -
In [70]: a = np.random.randint(0,9,(2,1,3,4))
In [71]: a
Out[71]:
array([[[[6, 4, 8, 6],
[4, 0, 1, 0],
[5, 3, 2, 5]]],
[[[3, 6, 0, 4],
[6, 2, 5, 2],
[0, 8, 0, 8]]]])
In [72]: sortidxs = np.argsort(np.linalg.norm(a, axis=-1))
In [73]: sortidxs
Out[73]:
array([[[1, 2, 0]],
[[0, 1, 2]]])
In [74]: a[np.arange(a.shape[0])[:,None,None],0, sortidxs]
Out[74]:
array([[[[4, 0, 1, 0],
[5, 3, 2, 5],
[6, 4, 8, 6]]],
[[[3, 6, 0, 4],
[6, 2, 5, 2],
[0, 8, 0, 8]]]])
答案 1 :(得分：1)
由于sortidxs包含每个轴的期望索引（从开始到结束），您可以通过np.arange(a.shape[0])生成第一个轴ragne并在索引时将其作为第一个轴传递：
In [31]: x,y, z = a.shape
In [32]: i, j = sortidxs.shape
In [33]: a[np.repeat(np.arange(x)[:, none], i, 1),sortidxs]
Out[33]:
array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])
作为一种更简化的方式，在这种情况下（正如您在评论中提到的那样），您只需传递np.arange(x)[:, None]而不使用repeat()函数，但如果您想要变体其他维数组repeat的第2和第3个索引或其他项目的数量将为您提供正确的答案。另请注意，在这种情况下，您还可以分别沿每个轴传递相应的索引。
In [107]: a[np.arange(x)[:, None],sortidxs]
Out[107]:
array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])