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

Question

设两个 ndarray:A 形状为 (n, *m) 和 B 形状为 (n, )。有没有办法使用对 B 进行排序的顺序来就地对 A 进行排序？

使用 np.argsort 对 A 和 B 进行排序很容易，但这不是就地完成的:

A = A[np.argsort(B)]

---

评论:

• A 和 B 具有不同的数据类型，并且 A 可以具有两个以上的维度。因此，它们不能堆叠起来使用 ndarray.sort()。
• A 占用大量空间，这就是为什么需要就地排序。因此，任何需要两倍于 A 占用空间的解决方案都将无法实现这一目的。
• 此问题的标题“Re-arranging numpy array in place ” 听起来可能相关，但问题本身不是很清楚，而且答案与我的问题不符。

Answer 1

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

为了保持 pythran 兼容，我必须稍微简化它，因此它只接受 2D 数组，并且只沿轴 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