Взятие вдоль оси в NumPy

Взять элемент в каждой строке по своему номеру позволяет функция take_along_axis. Она принимает массив, массив индексов той же числа размерностей и номер оси. При axis=1 в каждой строке выбирают свой столбец. Индексы задают не одну ячейку на весь массив, а свой столбец для каждой строки.

Из таблицы 2×3 возьмём последний столбец первой строки и первый столбец второй:

import numpy as np

table = np.array([[10, 20, 30], [40, 50, 60]])
cols = np.array([[2], [0]])
part = np.take_along_axis(table, cols, axis=1)
print(part) # выведет [[30] [40]]

Из таблицы [[1, 2], [3, 4], [5, 6]] возьмите по одному элементу из каждой строки: индексы 0, 1, 0 вдоль оси столбцов и выведите столбец.

Из таблицы [[9, 8, 7], [6, 5, 4]] возьмите элементы с индексами [[2], [1]] вдоль оси 1 и выведите форму (2, 1).