CODEWITHSUNDEEP
pythonarraysnumpy
econra2017
econra2017

Reputation: 67

Sorting a 3D array by a given row

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.

Upvotes: 3

Views: 110

Answers (3)

bb1
bb1

Reputation: 7863

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

Upvotes: 1

thomasreynolds4881
thomasreynolds4881

Reputation: 93

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)

Upvotes: 2

Sash Sinha
Sash Sinha

Reputation: 22370

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

Upvotes: 2

Related Questions