Нарезать наборы столбцов в NumPy

Question

Рассмотрим numpy array как таковое:

>>> a = np.array([[1, 2, 3, 0, 1], [2, 3, 2, 2, 2], [0, 3, 3, 2, 2]])
>>> a
array([[1, 2, 3, 0, 1],
       [2, 3, 2, 2, 2],
       [0, 3, 3, 2, 2]])

И array, который содержит пары индексов столбцов для среза (определенный столбец может появляться в нескольких парах):

b = [[0,1], [0,3], [1,4]]

Как я могу нарезать / транслировать / шагать a, используя b, чтобы получить такой результат:

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

       [[1, 0],
        [2, 2],
        [0, 2]],

       [[2, 1],
        [3, 2],
        [3, 2]]])

Answer 1

Используйте b в качестве индексов столбцов для подмножества массива и затем транспонируйте результат:

a[:, b].swapaxes(0, 1)

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