Sorting a 3D array by a given row

Question

I have a 2x3x4 Numpy array that looks like this:

[[[1 3 2 4]
  [0 1 0 1]
  [0 1 2 3]]
 [[1 4 2 3]
  [1 0 1 1]
  [0 1 2 3]]]

How can I sort this via the first row in each 2D matrix, i.e. achieve the output:

[[[1 2 3 4]
  [0 0 1 1]
  [0 2 1 3]]
 [[1 2 3 4]
  [1 1 1 0]
  [0 2 3 1]]]

Note how the 2nd and 3rd rows in each 2D matrix also move indices along with the 1st row.

Answer 1

Maybe you could use a list comprehension:

import numpy as np

arr = np.array([[
    [1, 3, 2, 4],
    [0, 1, 0, 1],
    [0, 1, 2, 3],
], [
    [1, 4, 2, 3],
    [1, 0, 1, 1],
    [0, 1, 2, 3],
]])
print('Before:')
print(arr)
sorted_arr = np.array([a[:, a[0].argsort()] for a in arr])
print('After:')
print(sorted_arr)

Output:

Before:
[[[1 3 2 4]
  [0 1 0 1]
  [0 1 2 3]]
 [[1 4 2 3]
  [1 0 1 1]
  [0 1 2 3]]]
After:
[[[1 2 3 4]
  [0 0 1 1]
  [0 2 1 3]]
 [[1 2 3 4]
  [1 1 1 0]
  [0 2 3 1]]]