Сортировка трехмерного массива по двумерному индексу

Question

У меня настоящие проблемы с этим. У меня есть трехмерный массив NumPy, и я хотел бы переупорядочить его с помощью двумерного индексного массива. В действительности массивы будут определяться программно, и трехмерный массив может быть двух или четырехмерным, но для простоты, вот желаемый результат, если оба массива были двумерными:

ph = np.array([[1,2,3], [3,2,1]])
ph_idx = np.array([[0,1,2], [2,1,0]])
for sub_dim_n, sub_dim_ph_idx in enumerate(ph_idx):
    ph[sub_dim_n] = ph[sub_dim_n][sub_dim_ph_idx]

Этопревращает массив ph в:

array([[1, 2, 3],
       [1, 2, 3]])

Это то, что я хотел бы. Может кто-нибудь помочь, если это то же самое обстоятельство, но вместо ph у меня есть трехмерный массив (psh), например:

psh = np.array(
    [[[1,2,3]], 
     [[3,2,1]]]
)

Надеюсь, что это понятно, и, пожалуйста, спросите, если это не так. Заранее спасибо!

Answer 1

Если вы хотите получить массив ph.shape, вы можете просто np.squeeze ph_ixs, чтобы формы соответствовали друг другу, и использовать его для индексации ph:

print(ph)
[[[1 2 3]]
 [[3 2 1]]]

print(ph_idx)
[[0 1 2]
 [2 1 0]]

np.take_along_axis(np.squeeze(ph), ph_idx, axis=-1)

array([[1, 2, 3],
       [1, 2, 3]])

Answer 2

Итак, подсказки уже есть в полезных комментариях, но для полноты информации это так же просто, как использовать np.take_along_axis и транслируемую версию массива 2d:

psh = np.array(
    [[[1,2,3]], 
     [[3,2,1]]]
)
ph_idx = np.array(
    [[0,1,2], 
     [2,1,0]]
)
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)

Это имеет преимуществотакже работает, если массив 3d имеет dim1 из более чем одного элемента:

psh = np.array(
    [[[1,2,3],[4,5,6],[7,8,9]], 
     [[3,2,1],[6,5,4],[9,8,7]]]
)
ph_idx = np.array([[0,1,2], [2,1,0]])
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)

, что дает

array([[[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]],

       [[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]])