Question

我有点困难。我试图在python中向量化一些代码以使其更快。我有一个数组，我排序（A）并获取索引列表（Ind）。我有另一个数组（B），我想按索引列表排序，而不使用我认为会阻碍计算的循环。

A = array([[2, 1, 9],
           [1, 1, 5],
           [7, 4, 1]])
Ind = np.argsort(A)

这是Ind的结果：

Ind = array([[1, 0, 2],
             [0, 1, 2],
             [2, 1, 0]], dtype=int64)

B是我想按Ind排序的数组：

B = array([[ 6,  3,  9],
           [ 1,  5,  3],
           [ 2,  7, 13]])

我想使用Ind重新排列B中的元素（B行按A行索引排序）：

B = array([[ 3,  6,  9],
           [ 1,  5,  3],
           [13,  7,  2]])

任何想法？我很乐意得到任何好的建议。我想提一下我使用的数百万个值，我的意思是30000 * 5000的数组。

干杯，罗伯特

Answer 1

我会做这样的事情：

import numpy as np
from numpy import array

A = array([[2, 1, 9],
           [1, 1, 5],
           [7, 4, 1]])
Ind = np.argsort(A)

B = array([[ 3,  6,  9],
           [ 1,  5,  3],
           [13,  7,  2]])

# an array of the same shape as A and B with row numbers for each element
rownums = np.tile(np.arange(3), (3, 1)).T

new_B = np.take(B, rownums * 3 + Ind)
print(new_B)
# [[ 6  3  9]
#  [ 1  5  3]
#  [ 2  7 13]]

您可以用数组形状替换幻数3。

通过Python中的另一个数组排序索引对数组进行排序

1 个答案: