Pythonic and efficient way to sort 2D array along column

Question

I have two arrays where, A is one dimensional. B is a 2D array (square matrix with same number of elements along both axes, as in A).

I needed to sort columns in B while sorting A in increasing order. The following code works but I wonder if it can be done without employing assignment to a new array and a for loop.

Looking for more efficient ways for get the same done. Thanks.

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) 

Output :

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

You can most certainly slice into B in a vectorized way. You can even do this without temporarily creating the index array:

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

More Detailed Output

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]])