python - 多维数组快速排序

Question

我需要根据第一个子数组中的值尽快对多维数组进行排序(该行被应用了数百万次)。

下面是我原来的路线，以及我尝试改善其性能的尝试，但没有成功。据我所知，我的 numpy 方法仅对第一个子数组进行正确排序，其余子数组均不排序。

我做错了什么以及如何提高排序的性能？

import numpy as np

# Generate some random data.
# I receive the actual data as a list, hence the .tolist()
aa = np.random.rand(10, 2000).tolist()

# This is the original line I need to process faster.
b1 = zip(*sorted(zip(*aa), key=lambda x: x[0]))

# This is my attempt at improving the above line's performance
b2 = np.sort(np.asarray(aa).T, axis=0).T

# Check if all sub-arrays are equal
for a, b in zip(*[b1, b2]):
    print(np.array_equal(a, b))

Answer 1

对于lambdas还是个新手，但是从我从你的代码中了解到的一点点来看 - 它似乎在你的 lambda 中方法，您正在使用x[0]获取排序键，然后使用它们从 aa 中的每个元素中提取值。用 NumPy 术语来说，这意味着获取数组版本中第一行的排序索引，然后索引到每一行(因为 aa 的每个元素成为数组 a 的每一行)。这基本上就是列索引。另外，似乎sorted保持相同元素的顺序。所以，我们需要使用argsort(kind='mergesort') .

因此，我们可以简单地做 -

a[:, a[0].argsort(kind='mergesort')] # a = np.array(aa) 

在您的 NumPy 代码中，您没有执行任何此类操作，因此没有给出正确的结果。