Блог

У меня есть трехмерный массив, который мне нужно отсортировать. У меня также есть 2-мерный индекс, который я получил с помощью a lexsort , который я хотел бы использовать для сортировки. Я не могу найти, как применить индекс к массиву без появления дополнительных измерений…

 import numpy as np

a = np.array([
    [
        [1, 1, 1],
        [0, 1, 2],
        [3, 1, 1],
        [-1, 1, 1],
        [0, 2, 2],
        [1, 2, 2],
    ],[
        [-1, 1, 1],
        [0, 1, 2],
        [-3, 1, 1],
        [1, 1, 1],
        [0, -2, 2],
        [-1, -2, 2],
    ],[
        [-1, 1, 1],
        [-0, 1, -2],
        [-3, 1, 1],
        [1, 1, 1],
        [-0, -2, -2],
        [-1, -2, -2],
    ],
])
sorted_index = np.lexsort((a[:,:,0], a[:,:,1],a[:,:,2]),axis=1)

  

sorted_index — это

 [[3 0 2 1 4 5], [2 0 3 5 4 1], [5 4 1 2 0 3]]
  

конечный результат должен быть

 a_sorted = np.array([
    [
        [-1, 1, 1],
        [1, 1, 1],
        [3, 1, 1],
        [0, 1, 2],
        [0, 2, 2],
        [1, 2, 2],
    ],[
        [-3, 1, 1],
        [-1, 1, 1],
        [1, 1, 1],
        [-1, -2, 2],
        [0, -2, 2],
        [0, 1, 2],
    ],[
        [-1, -2, -2],
        [0, -2, -2],
        [0, 1, -2],  
        [-3, 1, 1],
        [-1, 1, 1],
        [1, 1, 1],
    ],
])

  

Я пытался.

 a[sorted_index]
a[sorted_index, :]
a[:, sorted_index]
  

Я собираюсь выполнить цикл, но я хотел бы сделать это правильно.

Вы можете использовать np.take_along_axis

 np.take_along_axis(a,sorted_index[...,None],axis=1)
  

результат

 array([[[-1,  1,  1],
        [ 1,  1,  1],
        [ 3,  1,  1],
        [ 0,  1,  2],
        [ 0,  2,  2],
        [ 1,  2,  2]],

       [[-3,  1,  1],
        [-1,  1,  1],
        [ 1,  1,  1],
        [-1, -2,  2],
        [ 0, -2,  2],
        [ 0,  1,  2]],

       [[-1, -2, -2],
        [ 0, -2, -2],
        [ 0,  1, -2],
        [-3,  1,  1],
        [-1,  1,  1],
        [ 1,  1,  1]]])