我如何"压缩排序"并行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

b[a.argsort()] 应该做的伎俩.

这是它的工作原理.首先,你需要找到一个排序的排列.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])

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(),跳过冗余dtype,使用默认排序键,使用默认字段名而不是指定(基于此示例).