设两个ndarrays:形状为(n, *m)
的A
,和形状为(n, )
<的
B
的
/。有没有一种方法可以使用对
B
进行排序的顺序对 A
就地 进行排序?
使用 B
对 A
进行排序很容易使用 np.argsort
,但这不是就地完成的:
A = A[np.argsort(B)]
评论:
A
和B
有不同的 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