Сортировать матрицу numpy вдоль оси, но сохранить значения вместе

Question

У меня есть матрица, где 2-я ось (z / глубина) - это информация, которая объединяется; Подумайте, где три значения для каждого пикселя go вместе. Мой [y, x, классификатор] вместо [R, G, B].

a = np.array([[[10,15,1], [5,10,3], [10,5,3]],
               [[5,25,2], [20,20,3], [5,15,1]]]

Я бы хотел отсортировать их по 1-й и 2-й оси, чтобы наименьшее значение оси 1 + 2 было слева вверху, а наибольшее - справа внизу. Однако, когда я делаю np.sort(a, axis = 0, kind='stable'), он перемещает только первый элемент и не сохраняет их в виде набора. Я также попробовал:

idx = np.argsort(a, axis=0)
sorted_a = a[idx]

Но это стало очень странным и, казалось, повторяло всю матрицу несколько раз с различным порядком.

Конечным результатом должно быть:

[[[5,10,3], [5,15,1], [5,25,2]],
[[10,5,3], [10,15,1],[20,20,3]]]

Я бы предпочел метод, который использует встроенные методы NumPy, чтобы я мог заменить мои текущие простые циклы сортировки, которые в итоге стали довольно громоздкими .

Answer 1

Вот векторизованный способ -

In [211]: a2D = a.reshape(-1,a.shape[-1])

In [212]: (a2D[np.lexsort(a2D.T[::-1])]).reshape(a.shape)
Out[212]: 
array([[[ 5, 10,  3],
        [ 5, 15,  1],
        [ 5, 25,  2]],

       [[10,  5,  3],
        [10, 15,  1],
        [20, 20,  3]]])

Объяснение: Суть решения будет с lexsort, который в качестве состояния документов сортирует сгруппированные элементы, сохраняя определенный приоритетный порядок. В нашем случае, порядок, являющийся самым левым элементом, является наиболее важным и так далее до конца каждой строки. Основываясь на группировке, мы сортируем и, следовательно, получаем новый порядок строк. Мы используем его на двумерном массиве. Итак, мы меняем форму до и после сортировки по лексам.

Answer 2

Вы можете попробовать sorted и определить ключ сортировки, используя operator.itemgetter:

import operator

a = np.array([[[10,15,1], [5,10,3], [10,5,3]],
               [[5,25,2], [20,20,3], [5,15,1]]])
shp = a.shape
np.array(sorted(a.reshape(-1, 3), key=operator.itemgetter(0, 1))).reshape(shp)

Вывод:

array([[[ 5, 10,  3],
        [ 5, 15,  1],
        [ 5, 25,  2]],

       [[10,  5,  3],
        [10, 15,  1],
        [20, 20,  3]]])