我正在使用一个多维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)
或更多这样的“空”尺寸。我也希望排序在这些情况下起作用。
2条答案
按热度按时间1cklez4t1#
可以使用
advanced-indexing
-样品运行-
性能进一步提升
我们可以用
np.einsum
优化sortidxs
的计算-我们来计时并验证这个想法-
更高维数组
对于
a
是一个4D
数组的额外情况,我们需要使用更多的数组进行索引。1]对于第一轴:使用
np.arange(a.shape[0])
,并在末尾添加两个新轴。2]对于第二轴:使用
np.arange(a.shape[0])
,并在末尾添加一个新轴。3]对于第三轴:使用
sortidxs
索引到这个。因此,我们将有:
singleton dim(dim with length=1)数组
作为一种特殊情况,假设输入数组的第二个轴已经是单例的,我们可以简单地使用
0
作为该轴,从而简化事情,如下所示-样品运行-
另一个示例运行了一个通用形状为
(2,3,4)
的数组,以使事情变得非常清楚-wi3ka0sx2#
由于
sortidxs
包含每个轴(从开始到结束)的所需索引,因此可以通过np.arange(a.shape[0])
生成第一个轴范围,并在索引时将其作为第一个轴传递:作为一种更简单的方式,在这种情况下(正如你在评论中提到的),你可以只传递
np.arange(x)[:, None]
而不使用repeat()
函数,但是如果你想要沿着第二和第三个索引的项目的变量数量,或者其他维度数组repeat
会给予你正确的答案。还要注意,在这些情况下,您也可以分别沿着每个轴传递相应的索引。