Доступ к элементам numy из массива индексов

Question

У меня есть массив a измерения n и массив b измерения n-1.Значения последней оси b соответствуют индексам значений, которые я хочу извлечь из массива a в массиве res измерения n-1.

Например, при n = 2:

a = np.array([[1,  2,  3,  4],
              [5,  6,  7,  8],
              [9, 10, 11, 12])
b = np.array([1,3,0])

Мне бы хотелось

res = [a[1], a[3], a[0]]

# i.e. res = [2, 8, 9]

Есть ли функция, которая делает это эффективно, с большим количеством измерений?Я знаю, что мог бы использовать for петли, но я надеюсь, что есть что-то более эффективное.

РЕДАКТИРОВАТЬ:

При n = 3 пусть a будет иметь форму (2,2,3).

Тогда b и res имеют формы (2,2):

a = np.array([[[ 1, 2, 3],
               [ 4, 5, 6]],
              [[ 7, 8, 9],
               [10,11,12]]])

b = np.array([[0,2],
              [1,2]]

# res = np.array([[1,6],
#                 [8,12]])

Answer 1

Последний numpy (1.15) добавил функцию take_along_axis:

In [36]: np.take_along_axis(a, b[:,None], 1)
Out[36]: 
array([[2],
       [8],
       [9]])

. Он использует вспомогательную функцию для создания кортежа индексации:

In [37]: np.lib.shape_base._make_along_axis_idx((3,4), b[:,None], 1)
Out[37]: 
(array([[0],
        [1],
        [2]]), 
 array([[1],
        [3],
        [0]]))

Доэто, я (и другие) рекомендовал бы:

In [38]: a[np.arange(3), b]
Out[38]: array([2, 8, 9])

, что по сути то же самое (за исключением дополнительного измерения).Как показывают take_along_axis документы, это было разработано так, чтобы принимать такие вещи, как результаты argsort вдоль оси.

для случая с более высокой размерностью:

In [39]: a1 = np.array([[[ 1, 2, 3],
    ...:                [ 4, 5, 6]],
    ...:               [[ 7, 8, 9],
    ...:                [10,11,12]]])
    ...: b1 = np.array([[0,2],
    ...:               [1,2]])               
In [40]: a1.shape
Out[40]: (2, 2, 3)
In [41]: b1.shape
Out[41]: (2, 2)

In [42]: np.take_along_axis(a1, b1[...,None], -1)
Out[42]: 
array([[[ 1],
        [ 6]],

       [[ 8],
        [12]]])

In [45]: np.lib.shape_base._make_along_axis_idx(a1.shape, b1[...,None], 2)
Out[45]: 
(array([[[0]],

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

        [[1],
         [2]]]))
In [46]: [i.shape for i in _]
Out[46]: [(2, 1, 1), (1, 2, 1), (2, 2, 1)]

Снова эквивалентноеИндексирование самостоятельно:

In [48]: a1[np.arange(2)[:,None], np.arange(2)[None,:], b1]
Out[48]: 
array([[ 1,  6],
       [ 8, 12]])

Как только вы поймете, что такое вещание массива и как оно применяется к индексации, концепции здесь не сложны.Но take_along_axis может облегчить их применение.В некотором смысле это расширение np.ix_.

In [50]: np.ix_(np.arange(2), np.arange(2), np.arange(3))
Out[50]: 
(array([[[0]],

        [[1]]]), array([[[0],
         [1]]]), array([[[0, 1, 2]]]))
In [51]: [i.shape for i in _]
Out[51]: [(2, 1, 1), (1, 2, 1), (1, 1, 3)]
In [55]: a1[(*np.ix_(np.arange(2), np.arange(2)),b1)]
Out[55]: 
array([[ 1,  6],
       [ 8, 12]])