如果我有两个并列的列表,想要根据第一个列表中的元素顺序进行排序,这非常简单:
>>> 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()]
应该就可以解决问题。
它的原理是:首先需要找到一个对数组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])
b
作为“辅助排序”的标志,例如当 a
中有重复元素时。请参见我的答案以获取详细信息。 - Alok Singhal以下是一种方法,它不会创建任何中间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,使用默认排序键,使用默认字段名称而不是指定(基于这个示例)。
和@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]))
np.sort
的工作方式类似于list.sort
,并且在测试中没有注意到,因为示例数组应该以单独或词典方式排序。现在我已经给出了更好的答案(结果是@doug的答案的简化版本)。 - Matthias Fripp我遇到了相同的问题,并想知道不同排序方法和相应重排另一个数组的性能表现。
我认为这里提到的解决方案列表很全面,但我还在疑惑性能如何。因此,我实现了所有算法并进行了性能比较。
使用 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)
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)
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))
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仍然是最快的。对于其他需要辅助列排序的算法,基于recarray的解决方案现在比基于lexsort的变体更优。
免责声明:不同的数据类型和数组数据随机性可能会产生不同的结果。我使用42作为种子。
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有一个特殊之处,就是你必须按相反的顺序指定键,即将主键放在第二位,将次要键放在第一位。在我的例子中,我想使用第二列作为主键进行排序,所以我将其列为第二个键;第一列仅解决平局,但它被列为第一个键)。
a
和b
一起排序)与单独对a
和b
进行排序的解决方案相同。 - Steve