Как индексировать ndarrays

Доступ к определённым/произвольным строкам и столбцам

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

>>> a = np.arange(30).reshape(2, 3, 5)
>>> a
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14]],

        [[15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29]]])
>>> a[0, 2, :]
array([10, 11, 12, 13, 14])
>>> a[0, :, 3]
array([ 3,  8, 13])

Обратите внимание, что результат операций индексирования может иметь форму, отличную от исходного объекта. Для сохранения исходных размеров после индексирования можно использовать newaxis. Для использования других подобных инструментов см. инструменты многомерного индексирования.

>>> a[0, :, 3].shape
(3,)
>>> a[0, :, 3, np.newaxis].shape
(3, 1)
>>> a[0, :, 3, np.newaxis, np.newaxis].shape
(3, 1, 1)

Переменные также могут быть использованы для индексирования:

>>> y = 0
>>> a[y, :, y+3]
array([ 3,  8, 13])

Обратитесь к обработке переменного числа индексов в программах, чтобы узнать, как использовать slice и Ellipsis в ваших переменных индексов.

>>> a = np.arange(24).reshape(2, 3, 4)
>>> a
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])
>>> a[..., 3]
array([[ 3,  7, 11],
       [15, 19, 23]])

>>> arr = np.arange(3*4).reshape(3, 4)
>>> arr
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> column_indices = [[1, 3], [0, 2], [2, 2]]
>>> np.arange(arr.shape[0])
array([0, 1, 2])
>>> row_indices = np.arange(arr.shape[0])[:, np.newaxis]
>>> row_indices
array([[0],
       [1],
       [2]])

>>> arr[row_indices, column_indices]
array([[ 1,  3],
       [ 4,  6],
       [10, 10]])

>>> a = np.arange(30).reshape(2, 3, 5)
>>> a
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14]],

        [[15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29]]])
>>> np.take(a, [2, 3], axis=2)
array([[[ 2,  3],
        [ 7,  8],
        [12, 13]],

        [[17, 18],
        [22, 23],
        [27, 28]]])
>>> np.take(a, [2], axis=1)
array([[[10, 11, 12, 13, 14]],

        [[25, 26, 27, 28, 29]]])

Создание подмножеств больших матриц

Используйте срезы и шаги, чтобы получить фрагменты большого массива:

>>> a = np.arange(100).reshape(10, 10)
>>> a
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
        [60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
        [70, 71, 72, 73, 74, 75, 76, 77, 78, 79],
        [80, 81, 82, 83, 84, 85, 86, 87, 88, 89],
        [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])
>>> a[2:5, 2:5]
array([[22, 23, 24],
       [32, 33, 34],
       [42, 43, 44]])
>>> a[2:5, 1:3]
array([[21, 22],
       [31, 32],
       [41, 42]])
>>> a[:5, :5]
array([[ 0,  1,  2,  3,  4],
       [10, 11, 12, 13, 14],
       [20, 21, 22, 23, 24],
       [30, 31, 32, 33, 34],
       [40, 41, 42, 43, 44]])

То же самое можно сделать с расширенным индексированием несколько более сложным способом. Помните, что расширенное индексирование создаёт копию:

>>> a[np.arange(5)[:, None], np.arange(5)[None, :]]
array([[ 0,  1,  2,  3,  4],
       [10, 11, 12, 13, 14],
       [20, 21, 22, 23, 24],
       [30, 31, 32, 33, 34],
       [40, 41, 42, 43, 44]])

Также можно использовать mgrid для генерации индексов:

>>> indices = np.mgrid[0:6:2]
>>> indices
array([0, 2, 4])
>>> a[:, indices]
array([[ 0,  2,  4],
       [10, 12, 14],
       [20, 22, 24],
       [30, 32, 34],
       [40, 42, 44],
       [50, 52, 54],
       [60, 62, 64],
       [70, 72, 74],
       [80, 82, 84],
       [90, 92, 94]])

Фильтрация значений

>>> z = np.array([[1, 2, 3, 0], [0, 0, 5, 3], [4, 6, 0, 0]])
>>> z
array([[1, 2, 3, 0],
       [0, 0, 5, 3],
       [4, 6, 0, 0]])
>>> np.nonzero(z)
(array([0, 0, 0, 1, 1, 2, 2]), array([0, 1, 2, 2, 3, 0, 1]))

>>> np.flatnonzero(z)
array([0, 1, 2, 6, 7, 8, 9])

>>> a = np.arange(30).reshape(2, 3, 5)
>>> indices = np.where(a % 2 == 0)
>>> indices
(array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]),
array([0, 0, 0, 1, 1, 2, 2, 2, 0, 0, 1, 1, 1, 2, 2]),
array([0, 2, 4, 1, 3, 0, 2, 4, 1, 3, 0, 2, 4, 1, 3]))
>>> a[indices]
array([ 0,  2,  4,  6,  8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28])

>>> a > 14
array([[[False, False, False, False, False],
        [False, False, False, False, False],
        [False, False, False, False, False]],

       [[ True,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True]]])
>>> a[a > 14]
array([15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29])

>>> p = np.arange(-10, 10).reshape(2, 2, 5)
>>> p
array([[[-10,  -9,  -8,  -7,  -6],
        [ -5,  -4,  -3,  -2,  -1]],

       [[  0,   1,   2,   3,   4],
        [  5,   6,   7,   8,   9]]])
>>> q = p < 0
>>> q
array([[[ True,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True]],

       [[False, False, False, False, False],
        [False, False, False, False, False]]])
>>> p[q] = 0
>>> p
array([[[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

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

Получение индексов максимальных/минимальных значений

Используйте argmax и argmin:

>>> a = np.arange(30).reshape(2, 3, 5)
>>> np.argmax(a)
29
>>> np.argmin(a)
0

Используйте ключевое слово axis для получения индексов максимальных и минимальных значений по определённой оси:

>>> np.argmax(a, axis=0)
array([[1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1]])
>>> np.argmax(a, axis=1)
array([[2, 2, 2, 2, 2],
       [2, 2, 2, 2, 2]])
>>> np.argmax(a, axis=2)
array([[4, 4, 4],
       [4, 4, 4]])

>>> np.argmin(a, axis=1)
array([[0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]])
>>> np.argmin(a, axis=2)
array([[0, 0, 0],
       [0, 0, 0]])

Установите keepdims на True, чтобы сохранить оси, которые уменьшаются в результате, как измерения с размером один:

>>> np.argmin(a, axis=2, keepdims=True)
array([[[0],
        [0],
        [0]],

       [[0],
        [0],
        [0]]])
>>> np.argmax(a, axis=1, keepdims=True)
array([[[2, 2, 2, 2, 2]],

       [[2, 2, 2, 2, 2]]])

Для получения индексов каждого максимального или минимального значения для каждого (N-1)-мерного массива в N-мерном массиве используйте reshape для преобразования массива в 2D массив, примените argmax или argmin по axis=1 и используйте unravel_index для восстановления индекса значений по каждому слайсу:

>>> x = np.arange(2*2*3).reshape(2, 2, 3) % 7  # 3D example array
>>> x
array([[[0, 1, 2],
        [3, 4, 5]],

       [[6, 0, 1],
        [2, 3, 4]]])
>>> x_2d = np.reshape(x, (x.shape[0], -1))
>>> indices_2d = np.argmax(x_2d, axis=1)
>>> indices_2d
array([5, 0])
>>> np.unravel_index(indices_2d, x.shape[1:])
(array([1, 0]), array([2, 0]))

Первый возвращаемый массив содержит индексы по оси 1 в исходном массиве, второй массив содержит индексы по оси 2. Наибольшее значение в x[0] равно x[0, 1, 2].

Эффективное индексирование одного и того же ndarray несколько раз

Следует помнить, что основное индексирование создаёт представления, а расширенное индексирование создаёт копии, что менее эффективно с точки зрения вычислений. Поэтому следует стараться использовать основное индексирование, где это возможно, вместо расширенного.

Дополнительные материалы для чтения

Упражнения 100 по NumPy Никола Рюжье предоставляют хорошее понимание того, как индексирование сочетается с другими операциями. Упражнения 6, 8, 10, 15, 16, 19, 20, 45, 59, 64, 65, 70, 71, 72, 76, 80, 81, 84, 87, 90, 93, 94 специально сфокусированы на индексировании.