CODEWITHSUNDEEP
pythonsortingmultidimensional-array
math_enthusiast
math_enthusiast

Reputation: 369

A specific type of sorting in a 3d numpy array

In a 3d numpy array (with p panels, each with r rows and c columns) I'd like to sort only on columns of a specific panel, such that the corresponding elements on the other panels rearrange themselves accordingly.

Unfortunately, I am not familiar with the jargon of different types of sorting. I'll clarify what I need through an example.

Take A as a 2*3*4 array

A = array([[[   9, 20, 30, 11],
            [ 100,  4, -1, 90]
            [  40, 15, -5, 34]],

           [[  0,  2, 3,  9],
            [ -1, 12, 6, -3]
            [  1, -5, 7,  2]]]),

After sort on the columns of the second panel:

A = array([[[ 100, 15, 30, 90],
            [   9, 20, -1, 34]
            [  40,  4, -5, 11]],

           [[ -1, -5, 3, -3],
            [  0,  2, 6,  2]
            [  1, 12, 7,  9]]])

As you can see, only the columns of the second panel are sorted (ascendingly) and the elements in the first panel are rearranged (but not sorted!) with their corresponding elements in the second panel.

Upvotes: 2

Views: 81

Answers (1)

unutbu
unutbu

Reputation: 879749

import numpy as np

A = np.array([[[   9, 20, 30, 11],
            [ 100,  4, -1, 90],
            [  40, 15, -5, 34]],

           [[  0,  2, 3,  9],
            [ -1, 12, 6, -3],
            [  1, -5, 7,  2]]])

I, J, K = np.ogrid[tuple(map(slice, A.shape))]
# I, J, K are the identity indices in the sense that (A == A[I, J, K]).all()
newJ = np.argsort(A[1], axis=0) # first axis of A[1] is second axis of A
print(A[I, newJ, K])

yields

[[[100  15  30  90]
  [  9  20  -1  34]
  [ 40   4  -5  11]]

 [[ -1  -5   3  -3]
  [  0   2   6   2]
  [  1  12   7   9]]]

Upvotes: 3

Related Questions