python - 如何优化numpy 2D排序数组创建？

我有一个要排序的NxM矩阵，称为coefficients：

import numpy
N = 10
M = 42
coefficients = numpy.random.uniform(size=(N, M))

我有一个名为order的数组，其中包含N个元素，该数组表示coefficients的行应位于的顺序：

order = numpy.random.choice(range(N), N, False)

我正在对coefficients进行排序：

coefficients = numpy.array([mag for (orig, mag)
                            in sorted(zip(order, coefficients),
                                      key=lambda pair: pair[0])])

这行得通，但它可能比预期的要慢。如果是1D模式，我会使用order，但由于是2D模式，所以我不知道如何解决。我可以在这里进行优化吗？

最佳答案

要回答您的问题，只需coefficients[order.argsort()]就足够了:)

另请参见Numpy: sort by key function。