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.
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]]]
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)
One more possible solution:
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]]])
r, s, t = arr.shape
x, y, _ = np.ogrid[:r, :s, :t]
z = arr[:, 0, :].argsort()[:, np.newaxis, :]
out = arr[(x, y, z)]
print(out)
It gives:
[[[1 2 3 4]
[0 0 1 1]
[0 2 1 3]]
[[1 2 3 4]
[1 1 1 0]
[0 2 3 1]]]
