如何对平行的numpy数组进行“压缩排序”？

Question

如果我有两个平行的列表，想要按照第一个列表的顺序来排序它们，这个操作非常简单：

>>> a = [2, 3, 1]
>>> b = [4, 6, 7]
>>> a, b = zip(*sorted(zip(a,b)))
>>> print a
(1, 2, 3)
>>> print b
(7, 4, 6)

那么我怎么才能用numpy数组来做到这一点，而不需要把它们拆分成普通的Python列表呢？

Answer 1

就像@Peter Hansen的回答一样，这段代码在排序之前会先复制数组。不过它的实现很简单，主要的排序是在原地进行的，第二个数组是用来辅助排序的，应该会非常快：

a = np.array([2, 3, 1])
b = np.array([4, 6, 2])
# combine, sort and break apart
a, b = np.sort(np.array([a, b]))

更新: 上面的代码实际上是有问题的，正如评论中提到的那样。下面是一些更好的代码。这段代码应该相当高效，比如它避免了显式地创建额外的数组副本。很难说它的效率到底有多高，因为文档没有提供关于numpy.lexsort算法的详细信息。不过它应该能很好地工作，因为这正是lexsort被设计用来完成的任务。

a = np.array([5, 3, 1])
b = np.array([4, 6, 7])
new_order = np.lexsort([b, a])
a = a[new_order]
b = b[new_order]
print(a, b)
# (array([1, 3, 5]), array([7, 6, 4]))

Answer 2

这里有一种方法，它不会创建中间的 Python 列表，但需要用到 NumPy 的“记录数组”来进行排序。如果你的两个输入数组实际上是相关的（就像电子表格中的列），那么这可能会为你处理数据提供一种更好的方式，而不是一直保持两个独立的数组。在这种情况下，你可能已经有了一个记录数组，只需对这个数组调用 sort() 方法，就能解决你最初的问题。

这个方法在将两个数组打包成记录数组后，会进行一种原地排序：

>>> from numpy import array, rec
>>> a = array([2, 3, 1])
>>> b = array([4, 6, 7])
>>> c = rec.fromarrays([a, b])
>>> c.sort()
>>> c.f1   # fromarrays adds field names beginning with f0 automatically
array([7, 4, 6])

已编辑，使用 rec.fromarrays() 来简化，跳过多余的数据类型，使用默认的排序关键字，使用默认的字段名称而不是指定（基于这个例子）。

Answer 3

b[a.argsort()] 这个方法可以解决问题。

它的原理是这样的。首先，你需要找出一个可以把数组 a 排序的顺序。argsort 这个方法就是用来计算这个顺序的：

>>> a = numpy.array([2, 3, 1])
>>> p = a.argsort()
>>> p
[2, 0, 1]

你可以很简单地验证这个结果是否正确：

>>> a[p]
array([1, 2, 3])

接下来，把同样的顺序应用到数组 b 上。

>>> b = numpy.array([4, 6, 7])
>>> b[p]
array([7, 4, 6])