Sorting a 3D array by a given row-CodePudding

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.

CodePudding user response：

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

CodePudding user response：

You can use the transpose function to flip the axes and then more easily sort.

# fill in array
my_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]]])

# flip axes
new_arr = my_arr.transpose(0, 2, 1)

# create a new sorted array
new_arr_sorted = []
for arr in new_arr:
    new_arr_sorted.append(sorted(arr, key=lambda arr: int(''.join(map(str, arr)))))

# convert back to np array and flip back axes
new_arr_sorted = np.array(new_arr_sorted)
new_arr_sorted = new_arr_sorted.transpose(0, 2, 1)