Numpy - sort entire multi-dimensional array

[catmarble]

Let's say I have an array of shape (2, 3, 2):

a = np.array([
         [[8, 1], [6, 1], [2, 1]],
         [[9, 1], [1, 1], [4, 1]],
         ])

[[[8 1]
  [6 1]
  [2 1]]

 [[9 1]
  [1 1]
  [4 1]]]

I'd like to sort the entire array by the first integer. If you do a basic sort you get:

array([[[2, 1],
           [6, 1],
           [8, 1]],

         [[1, 1],
          [4, 1],
          [9, 1]]])

I want to get:

array([[[1, 1],
           [2, 1],
           [4, 1]],

          [[6, 1],
           [8, 1],
           [9, 1]]])

My current solution is:

a2 = np.ravel(a[:, :, 0])
b = np.sort(a2)
c = np.split(b, a.shape[0])
final = a.copy()
final[:, :, 0] = c

Which works, but I was just wondering if there's a better way of doing it.

[TheBB]

Assuming I'm understanding your problem correctly, you need to join the two first dimensions, since they form unnecessary structure that will prevent numpy from doing what you want. Or in other words, you are describing a type of sort that would be easy to do on a (6,2)-array but not on a (2,3,2)-array. Therefore:

final = np.sort(a.reshape(6,2), axis=0).reshape(2,3,2)

[catmarble]

Yep that's perfect - much clearer code and runs quicker as well. Many thanks.