Взятие различных столбцов из каждого 2D-среза массива 3D numpy

Question

Предположим следующий массив 3D numpy:

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

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

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

При заданном списке [3, 4, 8],

можно ли нарезать данный тензор без используя a для l oop?

Например, чтобы взять 3-й столбец из [0, :, :], 4-й столбец из [1, :, :] и 8-й столбец из [2, :, :], чтобы получить:

array([[5, 2, 2, 0, 1],
       [3, 5, 1, 5, 4],
       [2, 1, 3, 0, 0]])

Answer 1

Вот один способ с np.take_along_axis -

In [73]: idx = np.array([3,4,8])

# a is input array
In [72]: np.take_along_axis(a,idx[:,None,None],axis=2)[:,:,0]
Out[72]: 
array([[5, 2, 2, 0, 1],
       [3, 5, 1, 5, 4],
       [2, 1, 3, 0, 0]])

Другой способ с явным целочисленным индексированием -

In [79]: a[np.arange(len(idx)),:,idx]
Out[79]: 
array([[5, 2, 2, 0, 1],
       [3, 5, 1, 5, 4],
       [2, 1, 3, 0, 0]])