NumPy: вычисления и устройство массивов

Агрегации по массиву и по оси: sum, mean, min, max, argmin, argmax

Содержание курса

Сквозной пример и процедура проверки shape

Теперь соберём всё вместе на одном массиве и отработаем конкретный алгоритм: сначала предсказать shape результата вручную, потом проверить через .shape.

Возьмём уже знакомый массив:

import numpy as np

arr = np.array([[3, 1, 4],
                [1, 5, 9],
                [2, 6, 5]])
# shape: (3, 3)

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

  • np.sum(arr) — axis не указан → все оси схлопываются → скаляр, shape ()
  • np.mean(arr, axis=0) — убираем элемент с позицией 0 из (3, 3)(3,)
  • np.min(arr, axis=1) — убираем элемент с позицией 1 из (3, 3)(3,)
  • np.argmax(arr) — axis не указан → позиция в flatten → скаляр, shape ()
  • np.argmin(arr, axis=1) — убираем элемент с позицией 1 из (3, 3)(3,)

Теперь проверяем:

# Глобальные агрегации → скаляр
print(np.sum(arr))             # 36
print(np.sum(arr).shape)       # ()

print(np.mean(arr))            # 4.0
print(np.min(arr))             # 1
print(np.max(arr))             # 9

# По оси 0 → shape (3,)
res_axis0 = np.mean(arr, axis=0)
print(res_axis0)               # [2.  4.  6.]
print(res_axis0.shape)         # (3,)

# По оси 1 → shape (3,)
res_axis1 = np.min(arr, axis=1)
print(res_axis1)               # [1 1 2]
print(res_axis1.shape)         # (3,)

# argmax без axis → скаляр (позиция в flatten)
print(np.argmax(arr))          # 5
print(np.argmax(arr).shape)    # ()

# argmin по оси 1 → shape (3,)
res_argmin = np.argmin(arr, axis=1)
print(res_argmin)              # [1 0 0]
print(res_argmin.shape)        # (3,)

Все предсказания совпадают. Это и есть рабочий алгоритм:

  1. Запишите shape исходного массива.
  2. Если axis не указан — результат всегда скаляр, shape ().
  3. Если указан axis=k — уберите из кортежа shape элемент на позиции k. Остаток и есть shape результата.
  4. Запустите агрегацию и вызовите .shape — сверьте с предсказанием.

Алгоритм работает одинаково для sum, mean, min, max, argmin, argmax — и для массивов любой размерности. Например, для трёхмерного:

arr3d = np.ones((2, 3, 4))
# axis=1 → убираем '3' из (2, 3, 4) → предсказание (2, 4)
result = np.sum(arr3d, axis=1)
print(result.shape)  # (2, 4) ✓

Привычка проверять .shape сразу после агрегации экономит время: несовпадение с предсказанием сразу указывает, что axis выбран неверно или массив имеет другую форму, чем ожидалось.