Numpy Array Permute Columns 3d Matrix With Argsort

I need to permute elements of columns in the matrix A (3D matrix by axis 0) by 2D permutation matrix pi obtained from argsort, that contains new indices for all columns. By applica

Solution 1:

Maybe a matter of taste, but I find the following a bit more readable:

i, j = np.ogrid[:3, :4]
A[pi[..., None], i, j]

Output:

array([[[1335,1747,3418,6312],
        [ 5515,  8292,  8414, 16135],
        [  565,  3038,  3800,  5430]],

       [[   73,   701,  2411,  2414],
        [ 3788,  5449,  5753,  9738],
        [  100,  1241,  2146,  2931]]])