Skip to main content
 首页 » 编程设计

python之使用第二个数组的顺序就地对 numpy 数组进行排序

2024年09月03日7mate10pro

设两个ndarrays:形状为(n, *m)A,和形状为(n, ) <的 B /。有没有一种方法可以使用对 B 进行排序的顺序对 A 就地 进行排序?

使用 BA 进行排序很容易使用 np.argsort ,但这不是就地完成的:

A = A[np.argsort(B)] 

评论:

  • AB 有不同的 dtype,A 可以有两个以上的维度。因此它们不能堆叠使用 ndarray.sort()
  • A 占用了大量空间,这就是为什么需要就地排序的原因。因此,任何需要 A 所占用空间两倍的解决方案都会破坏这一目的。
  • 这个问题的标题“Re-arranging numpy array in place ” 听起来可能相关,但问题本身不是很清楚,答案与我的问题不符。

请您参考如下方法:

这是一个通过在索引数组中跟踪循环来工作的解决方案。它可以选择使用 pythran 进行编译,如果行很小(10 个元素为 80 倍),则可以提供显着的加速,如果行很大(1000 个元素,则为 30%)。

为了让它与 pythran 兼容,我不得不稍微简化它,所以它只接受二维数组,并且只沿轴 0 排序。

代码:

import numpy as np 
 
#pythran export take_inplace(float[:, :] or int[:, :], int[:]) 
 
def take_inplace(a, idx): 
    n, m = a.shape 
    been_there = np.zeros(n, bool) 
    keep = np.empty(m, a.dtype) 
    for i in range(n): 
        if been_there[i]: 
            continue 
        keep[:] = a[i] 
        been_there[i] = True 
        j = i 
        k = idx[i] 
        while not been_there[k]: 
            a[j] = a[k] 
            been_there[k] = True 
            j = k 
            k = idx[k] 
        a[j] = keep 

使用编译版本运行示例。如上所述,只有小行需要编译,对于大行,纯 Python 应该足够快。

>>> from timeit import timeit 
>>> import numpy as np 
>>> import take_inplace 
>>>  
>>> a = np.random.random((1000, 10)) 
>>> idx = a[:, 4].argsort() 
>>> 
>>> take_inplace.take_inplace(a, idx) 
>>> 
# correct 
>>> np.all(np.arange(1000) == a[:, 4].argsort()) 
True 
>>> 
# speed 
>>> timeit(lambda: take_inplace.take_inplace(a, idx), number=1000) 
0.011950935004279017 
>>> 
# for comparison 
>>> timeit(lambda: a[idx], number=1000) 
0.02985276997787878