Многомерная индексация — NumPy

Коротко

Индексы по осям пишутся через запятую в одних скобках: a[строка, столбец].
a[i] — строка i; a[:, j] — столбец j; a[1:3, 2:] — блок.
Целый индекс убирает ось, срез её сохраняет: a[:, 2] имеет форму (n,), а a[:, 2:3] — (n, 1).
... (Ellipsis) заменяет «все остальные оси»: a[..., 0] — последний индекс равен 0.
Как и в одномерном случае, срезы — представления: запись в них меняет исходный массив.

Синтаксис

a[i, j]           # один элемент
a[i]              # строка i (то же, что a[i, :])
a[:, j]           # столбец j
a[r1:r2, c1:c2]   # блок
a[..., k]         # последний индекс k, остальные — все
a[i, :] = 0       # запись в строку

Примеры

Данные для примеровиспользуются во всех примерах ниже

# температура: 3 города (строки) × 5 дней (столбцы)
temps = np.array([[12, 14, 15, 13, 11],
                  [ 8,  9, 11, 10,  7],
                  [20, 22, 21, 19, 18]])

Один элемент

print(temps[1, 2])      # город 1, день 2
temps[-1, -1]           # последний город, последний день

Вывод

11
np.int64(18)

Строка и столбец

print(temps[0])         # строка — все дни города 0
temps[:, 2]             # столбец — все города в день 2

Вывод

[12 14 15 13 11]
array([15, 11, 21])

Блок и шаг

print(temps[:2, 1:4])   # два первых города, дни 1–3
temps[:, ::2]           # каждый второй день

Вывод

[[14 15 13]
 [ 9 11 10]]
array([[12, 15, 11],
       [ 8, 11,  7],
       [20, 21, 18]])

Многоточие для многомерных массивов

cube = np.arange(24).reshape(2, 3, 4)
print(cube[..., 0])     # то же, что cube[:, :, 0]
cube[0, ...].shape

Вывод

[[ 0  4  8]
 [12 16 20]]
(3, 4)

Запись в строку и столбец

temps[1, :] = 0          # обнулить строку
temps[:, 0] = [1, 2, 3]  # записать столбец
temps

Вывод

array([[ 1, 14, 15, 13, 11],
       [ 2,  0,  0,  0,  0],
       [ 3, 22, 21, 19, 18]])

Подводные камни

Столбец теряет ось

a[:, 2] — это вектор (n,), а не столбец (n, 1). Для транслирования и склейки столбцов часто нужна именно форма (n, 1) — берите срез длины 1 или добавьте ось через np.newaxis.

Целый индекс и срез длины 1

print(temps[:, 2].shape)     # столбец превратился в вектор
temps[:, 2:3].shape          # срез длины 1 сохраняет ось

Вывод

(3,)
(3, 1)

a[:][0] — не столбец

Две пары скобок применяются по очереди: a[:] возвращает весь массив, а [0] берёт из него первую строку. Всегда пишите индексы через запятую: a[:, 0]. Это ещё и быстрее.

Цепочка скобок

print(temps[:, 0])      # столбец 0
temps[:][0]             # это НЕ столбец: [:] — все строки, затем [0] — первая строка

Вывод

[12  8 20]
array([12, 14, 15, 13, 11])