Могу ли я «вращать» и нарезать за одну операцию?

Question

У меня есть такие данные:

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

и мне нужно привести его в эту "форму":

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

Лучший способ сделать это на данный момент:

for i in range(2):
    pairs = tuple(
        array[i, :, j] for j in range(3)  # ---axis-2-length---
    )
    print(pairs)

произвести:

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

Хотя это и делает свою работу, это очень много интерпретируется Python, особенно с ростом этих наборов данных. Так что мне было интересно, смогу ли я как-нибудь сделать это с помощью причудливого индексирования или транспонирования, или еще чего-нибудь, но моя глупая фу еще не достаточно сильна, чтобы представить себя.

Кто-нибудь знает, как сделать это более элегантно с NumPy?

Answer 1

Транспонирование с ndarray.transpose будет работать здесь.

>>> x.transpose(0, 2, 1)
array([[[1, 4],
        [2, 5],
        [3, 6]],

       [[7, 0],
        [8, 1],
        [9, 2]]])