沿列对 2D 数组进行排序的 Pythonic 且有效的方法

Question

我有两个数组，其中 A 是一维的。 B 是一个二维数组(沿两个轴具有相同数量元素的方阵，与 A 中相同)。

我需要对 B 中的列进行排序，同时对 A 进行升序排序。 以下代码有效，但我想知道是否可以在不使用分配给新数组和 for 循环的情况下完成。

寻找更有效的方法来完成同样的任务。谢谢。

import numpy as np

A=np.random.randint(1,50,6)
B=np.random.randint(1,50,(6,6))
print ("A = ",A,"\n \nB = \n", B)

ind = np.argsort(A, axis=0)
print("\nSorting index = \n",ind)

C=A[ind]
print("\nSorted A = \n",C)

D=np.empty_like(B)
for i in range(6):
    D[:,i] = B[:,ind[i]]

print("\nSorted B along the columns = \n",D) 

输出:

A =  [40  1 12 42 15  3] 

B = 
 [[43 20 26 15 24 13]
 [36  7 47 14 36 11]
 [44 19 41 32 14 43]
 [27 11 46 44 35 22]
 [26 18  4 40 40 23]
 [27 23 30 49 28 12]]

Sorting index = 
 [1 5 2 4 0 3]

Sorted A = 
 [ 1  3 12 15 40 42]

Sorted B along the columns = 
 [[20 13 26 24 43 15]
 [ 7 11 47 36 36 14]
 [19 43 41 14 44 32]
 [11 22 46 35 27 44]
 [18 23  4 40 26 40]
 [23 12 30 28 27 49]]

Answer 1

您肯定可以以矢量化方式切片到 B 中。您甚至可以在不临时创建索引数组的情况下执行此操作:

D = B[:,np.argsort(A, axis=0)]

更详细的输出

In [13]: A
Out[13]: array([40,  1, 12, 42, 15,  3])

In [14]: B
Out[14]:
array([[43, 20, 26, 15, 24, 13],
       [36,  7, 47, 14, 36, 11],
       [44, 19, 41, 32, 14, 43],
       [27, 11, 46, 44, 35, 22],
       [26, 18,  4, 40, 40, 23],
       [27, 23, 30, 49, 28, 12]])

In [15]: D = B[:,np.argsort(A, axis=0)]

In [16]: D
Out[16]:
array([[20, 13, 26, 24, 43, 15],
       [ 7, 11, 47, 36, 36, 14],
       [19, 43, 41, 14, 44, 32],
       [11, 22, 46, 35, 27, 44],
       [18, 23,  4, 40, 26, 40],
       [23, 12, 30, 28, 27, 49]])