numpy.take_along_axis#

numpy.take_along_axis(arr, индексы, ось=-1)[источник]#

Берут значения из входного массива путем сопоставления 1D индексов и срезов данных.

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

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

Параметры:
arrndarray (Ni..., M, Nk...)

Исходный массив

индексыndarray (Ni…, J, Nk…)

Тип данных файла. arr. Это должно соответствовать размерности arr, но размерности Ni и Nj должны только транслироваться на arr.

осьint или None, опционально

Ось, вдоль которой брать 1D срезы. Если axis равен None, входной массив обрабатывается так, как если бы он сначала был сведён к 1D, для согласованности с sort и argsort.

Изменено в версии 2.3: Значение по умолчанию теперь -1.

Возвращает:
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(Ni + (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]

Примеры

>>> import numpy as np

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

>>> 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]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])

То же работает для max и min, если сохранить тривиальную размерность с помощью keepdims:

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

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

>>> ai_min = np.argmin(a, axis=1, keepdims=True)
>>> ai_max = np.argmax(a, axis=1, keepdims=True)
>>> ai = np.concatenate([ai_min, ai_max], axis=1)
>>> ai
array([[0, 1],
       [1, 0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 30],
       [40, 60]])
На этой странице