使用第二个数组的顺序对 numpy 数组 in-place 进行排序

Question

设两个ndarrays：形状(n, *m)的A和形状(n, )的B。有没有一种方法可以使用排序 B?

的顺序对 A in-place 进行排序

用 B 对 A 进行排序很容易使用 np.argsort，但这还没有完成 in-place:

A = A[np.argsort(B)]

评论：

A 和 B 有不同的数据类型，A 可以有两个以上的维度。因此它们不能堆叠使用 ndarray.sort().
A占用了很多space，这就是为什么需要对它进行排序in-place。因此，任何需要 A 占用的 space 两倍的解决方案都会破坏此目的。
这个问题的标题“Re-arranging numpy array in place”可能听起来相关，但问题本身不是很清楚，答案与我的问题不符。

Answer 1

如果您可以预先将A设置为结构化数组，其数据类型由形状为(m, )的子数组和相同类型的标量组成（例如 np.int32），然后您可以根据 B 就地对其进行排序。例如：

import numpy as np

B = np.array([3, 1, 2])
A = np.array([[10, 11], [20, 21], [30, 31]])

(n, m) = A.shape

dt = np.dtype([('a', np.int32, (m, )), ('b', int)])
A2 = np.array([(a, b) for a, b in zip(A, B)], dtype=dt)

A2.sort(order='b')
print A2

Answer 2

这是一个通过在索引数组中跟踪循环来工作的解决方案。它可以选择使用 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

使用第二个数组的顺序对 numpy 数组 in-place 进行排序

Sort numpy array in-place using order from a second array

python

numpy

numpy-ndarray