如何“压缩排序”并行数组?

如果我有两个并行列表,并且希望按照第一个列表中元素的顺序对它们进行排序,这很简单:

>>> 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 列表中?

41432 次浏览

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“ record array”来进行排序。如果您的两个输入数组实际上是相关的(就像电子表格中的列) ,那么这可能会开辟一种处理数据的有利方式,而不是一直保持两个不同的数组,在这种情况下,您已经有了一个记录数组,并且您的原始问题将仅仅通过在数组上调用 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.frmarray () ,跳过冗余的 dtype,使用默认排序键,使用默认字段名而不是指定(基于 这个例子)。

这可能是做您想要做的事情的最简单和最一般的方法。(我在这里使用了三个数组,但是它可以用于任何形状的数组,无论是两列还是两百列)。

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 的一个怪异之处在于,您必须以相反的顺序指定键,即,将主键放在第二位,辅键放在第一位。在我的示例中,我想使用第2列作为主键进行排序,所以我将其列在第二列; 第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]))

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

性能比较两种数组情况

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

使用 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 重排进行排序

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 重排进行排序

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

对包含100个数组、每个数组包含100000个随机整数(arrays = [np.random.randint(10, size=100000) for _ in range (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 仍然是最快的,这似乎是合乎逻辑的,因为忽略了辅助排序。对于其他算法,那些使用辅助列排序的算法,基于重排的解决方案现在打败了 lexsort 变体。

免责声明: 其他 dtype 的结果可能不同,也取决于数组数据的随机性。我用42号做种子。