如何“压缩排序”并行 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()] 应该可以解决问题。

这是它的工作原理。首先，您需要找到对 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])

Answer 2

这是一种不创建中间 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）。

Answer 3

就像@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]))

Answer 4

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

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)

使用 lexsort 的一个怪癖是您必须以相反的顺序指定键，即，将主键放在第二位，将辅助键放在第一位。在我的示例中，我想使用第二列作为主键进行排序，因此我将其列为第二个；第 1 列仅解析平局，但首先列出）。

Answer 5

我遇到了同样的问题，想知道对一个数组进行排序并相应地重新排序另一个数组的不同方法的性能。

两个数组的性能比较案例

我认为这里提到的解决方案列表很全面，但我也想知道性能。因此我实现了所有算法并进行了性能比较。

使用 zip 排序两次

def zip_sort(s, p):
    ordered_s, ordered_p = zip(*sorted(list(zip(s, p))))
    return np.array(ordered_s, dtype=s.dtype), np.array(ordered_p, dtype=p.dtype)

使用 argsort 进行排序。这里不会考虑其他数组进行辅助排序

def argsort(s, p):
    indexes = s.argsort()
    return s[indexes], p[indexes]

使用 numpy recarrays 排序

def recarray_sort(s, p):
    rec = np.rec.fromarrays([s, p])
    rec.sort()
    return rec.f0, rec.f1

使用 numpy lexsort 进行排序

def lexsort(s, p):
    indexes = np.lexsort([p, s])
    return s[indexes], p[indexes]

对 100000 个随机整数的两个列表 p 和 q 进行排序将产生以下性能

zip_sort
258 ms ± 7.32 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

argsort
9.67 ms ± 278 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

recarray_sort
86.4 ms ± 707 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

lexsort
12.4 ms ± 288 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

因此 argsort 是最快的，但也会产生与其他算法略有不同的结果。如果不需要辅助排序，则应使用 argsort。

性能比较多数组案例

接下来，可能需要对多个数组进行这种排序。修改算法以处理多个数组看起来像

使用 zip 排序两次

def zip_sort(*arrays):
    ordered_lists = zip(*sorted(list(zip(*arrays))))
    return tuple(
        (np.array(l, dtype=arrays[i].dtype) for i, l in enumerate(ordered_lists))
    )

使用 argsort 进行排序。这里不会考虑其他数组进行辅助排序

def argsort(*arrays):
    indexes = arrays[0].argsort()
    return tuple((a[indexes] for a in arrays))

使用 numpy recarrays 排序

def recarray_sort(*arrays):
    rec = np.rec.fromarrays(arrays)
    rec.sort()
    return tuple((getattr(rec, field) for field in rec.dtype.names))

使用 numpy lexsort 进行排序

def lexsort(*arrays):
    indexes = np.lexsort(arrays[::-1])
    return tuple((a[indexes] for a in arrays))

用每 100000 个随机整数 (arrays = [np.random.randint(10, size=100000) for _ in range (100)]) 对 100 个数组的列表进行排序，现在可以得到以下性能

zip_sort
13.9 s ± 570 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

argsort
49.8 ms ± 1.49 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

recarray_sort
491 ms ± 14.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

lexsort
881 ms ± 15.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

argsort 仍然是最快的，由于忽略了辅助排序，这似乎是合乎逻辑的。对于其他算法，那些具有辅助列排序的算法，基于 recarray 的解决方案现在胜过 lexsort 变体。

免责声明：其他数据类型的结果可能会有所不同，并且还取决于数组数据的随机性。我用 42 作为种子。