Эффективное индексирование матрицы по столбцам в numpy

Question

У меня есть две матрицы одинакового размера, A, B. Я хочу использовать столбцы B, чтобы получить доступ к столбцам A для каждого столбца. Например,

A = np.array([[1, 4, 7],
              [2, 5, 8],
              [3, 6, 9]])

и

B = np.array([[0, 0, 2],
              [1, 2, 1],
              [2, 1, 0]])

Я хочу что-то вроде:

A[B] = [[1, 4, 9],  
        [2, 6, 8], 
        [3, 5, 7]]

Т.е. я использовал j-й столбец B как указатели к j-й колонке А. Есть ли эффективный способ сделать это? Спасибо!

Answer 1

Вы можете использовать расширенное индексирование :

A[B, np.arange(A.shape[0])]

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

Или с np.take_along_axis:

np.take_along_axis(A, B, axis=0)

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