我如何“压缩排序”并行numpy数组？

Posted 2021-04-02

如果我有两个并行列表，并希望按照第一个元素的顺序对它们进行排序，那么很容易：

>>> 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列表中呢？

答案

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])

另一答案

这是一种不创建中间Python列表的方法，但它确实需要NumPy“记录数组”来用于排序。如果你的两个输入数组实际上是相关的（比如电子表格中的列）那么这可能会打开一种处理数据的有利方式，而不是一直保持两个不同的数组，在这种情况下你已经有了记录数组，只需在数组上调用sort（）即可回答原始问题。

在将两个数组打包到记录数组中之后，这会执行in-place 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，使用默认排序键，使用默认字段名而不是指定（基于this example）。

另一答案

这可能是最简单，最通用的方式来做你想要的。 （我在这里使用了三个数组，但这适用于任何形状的数组，无论是两列还是两百个）。

import numpy as NP
fnx = lambda : NP.random.randint(0, 10, 6)
a, b, c = fnx(), fnx(), fnx()
abc = NP.column_stack((a, b, c))
keys = (abc[:,0], abc[:,1])          # sort on 2nd column, resolve ties using 1st col
indices = NP.lexsort(keys)        # create index array
ab_sorted = NP.take(abc, indices, axis=0)

一个怪癖w / lexsort是你必须以相反的顺序指定密钥，即，将主键放在第二位，将第二个键放在第一位。在我的例子中，我想使用第二列作为主键进行排序，所以我将其列为第二列;第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]))

更新：正如评论中所指出的，上面的代码实际上并不起作用。下面是一些更好的代码。这应该是相当有效的 - 例如，它避免明确地制作阵列的额外副本。很难说它会有多高效，因为the documentation没有提供有关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]))