numpy.take_along_axis

numpy.take_along_axis(arr, indices, axis) [source]

Получение значений из массива-источника по соответствию одномерных индексов и срезов данных.

Эта функция итерируется по соответствующим одномерным срезам, ориентированным вдоль указанной оси в массивах индексов и данных, и использует первый для поиска значений во втором. Эти срезы могут иметь разную длину.

Функции, возвращающие индекс вдоль оси, такие как argsort и argpartition, генерируют подходящие индексы для этой функции.

Добавлена в версии 1.15.0.

Параметры:

arr: ndarray (Ni…, M, Nk…) Массив-источник. indices: ndarray (Ni…, J, Nk…) Индексы для извлечения вдоль каждого одномерного среза массива arr. Они должны соответствовать размерности массива arr, но размерности Ni и Nj должны быть совместимы для трансляции с arr. axis: int Ось для извлечения одномерных срезов. Если axis равно None, массив-источник обрабатывается как если бы он был сперва сжат в одномерный массив, для согласованности с sort и argsort.

Возвращаемое значение:

out: ndarray (Ni…, J, Nk…) Результат извлечения по индексам.

См. также

take

Извлечение вдоль оси, используя одни и те же индексы для каждого одномерного среза

put_along_axis

Установка значений в целевой массив по соответствию одномерных индексов и срезов данных

Примечания

Это эквивалентно (но быстрее, чем) следующему использованию ndindex и s_, которое устанавливает каждый из ii и kk в кортеж индексов:

Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]
J = indices.shape[axis]  # Need not equal M
out = np.empty(Nk + (J,) + Nk)

for ii in ndindex(Ni):
    for kk in ndindex(Nk):
        a_1d       = a      [ii + s_[:,] + kk]
        indices_1d = indices[ii + s_[:,] + kk]
        out_1d     = out    [ii + s_[:,] + kk]
        for j in range(J):
            out_1d[j] = a_1d[indices_1d[j]]

В эквивалентном варианте, исключая внутренний цикл, две последние строки будут:

out_1d[:] = a_1d[indices_1d]

Примеры

Для этого массива

>>> a = np.array([[10, 30, 20], [60, 40, 50]])

Мы можем отсортировать, либо используя метод sort напрямую, либо argsort и эту функцию

>>> np.sort(a, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])
>>> ai = np.argsort(a, axis=1); ai
array([[0, 2, 1],
       [1, 2, 0]], dtype=int64)
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])

То же самое работает для max и min, если вы расширите размерности:

>>> np.expand_dims(np.max(a, axis=1), axis=1)
array([[30],
       [60]])
>>> ai = np.expand_dims(np.argmax(a, axis=1), axis=1)
>>> ai
array([[1],
       [0], dtype=int64)
>>> np.take_along_axis(a, ai, axis=1)
array([[30],
       [60]])

Если мы хотим получить max и min одновременно, мы можем сначала склеить индексы

>>> ai_min = np.expand_dims(np.argmin(a, axis=1), axis=1)
>>> ai_max = np.expand_dims(np.argmax(a, axis=1), axis=1)
>>> ai = np.concatenate([ai_min, ai_max], axis=axis)
>> ai
array([[0, 1],
       [1, 0]], dtype=int64)
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 30],
       [40, 60]])