Вычисления без циклов

Главная причина использовать numpy — скорость и краткость: вместо цикла по элементам пишется одна операция над всем массивом. Это называется векторизацией.

Почему без циклов быстрее

import numpy as np

nums = list(range(1_000_000))
squares = [x * x for x in nums]       # цикл Python: ~0.1 с

a = np.arange(1_000_000)
squares = a * a                        # numpy: ~0.002 с

Цикл Python на каждом шаге проверяет типы и создаёт объекты. numpy делает ту же работу внутри одного вызова на C. Разница — в десятки раз, а на больших данных это минуты против секунд.

Арифметика и функции поэлементно

Все операции применяются к каждому элементу:

a = np.array([1, 4, 9])
b = np.array([10, 20, 30])
print(a + b)          # [11 24 39]
print(a * b)          # [ 10  80 270]
print(a ** 2)         # [ 1 16 81]
print(np.sqrt(a))     # [1. 2. 3.]
print(np.abs(np.array([-2, 3])))   # [2 3]

У numpy есть свои версии математических функций: np.sqrt, np.exp, np.log, np.sin, np.round, np.abs. Обычный math.sqrt с массивом работать не будет.

Агрегаты: одно число из массива

t = np.array([12, 15, 9, 20, 14])
print(t.sum())        # 70
print(t.mean())       # 14.0 — среднее
print(t.min(), t.max())        # 9 20
print(t.argmin(), t.argmax())  # 2 3 — индексы минимума и максимума
print(t.std())        # стандартное отклонение

argmin и argmax возвращают позицию (с нуля). При нескольких одинаковых максимумах — позицию первого.

Ось axis

У таблицы агрегат можно считать по всей таблице, по столбцам или по строкам:

sales = np.array([[5, 3, 8],     # магазин 1 по трём дням
                  [2, 7, 4]])    # магазин 2
print(sales.sum())          # 29 — всё вместе
print(sales.sum(axis=0))    # [ 7 10 12] — по каждому дню (столбцу)
print(sales.sum(axis=1))    # [16 13] — по каждому магазину (строке)

Как запомнить: axis=0 — ось строк «схлопывается», результат — по одному числу на столбец. axis=1 — схлопываются столбцы, результат — по одному числу на строку.

Это работает для sum, mean, min, max, argmax и других агрегатов.

Broadcasting: разные размеры

Если размеры массивов не совпадают, numpy пытается «растянуть» меньший:

m = np.array([[1, 2, 3],
              [4, 5, 6]])
print(m * 10)                       # число растягивается на все элементы
print(m - np.array([1, 1, 1]))      # строка вычитается из каждой строки

Строка длины 3 подходит к таблице с тремя столбцами. А чтобы вычесть из каждой строки своё число, нужен столбец формы (2, 1):

row_means = m.mean(axis=1)              # [2. 5.]
centered = m - row_means.reshape(2, 1)  # из каждой строки — её среднее
print(centered)
# [[-1.  0.  1.]
#  [-1.  0.  1.]]

Если размеры совсем несовместимы (например, 3 и 4), будет ошибка operands could not be broadcast together.

Чтение таблицы

Таблицу из входных данных удобно читать сразу в массив:

n, m = map(int, input().split())
a = np.array([input().split() for _ in range(n)], dtype=float)

Список строк-частей превращается в таблицу n × m, а dtype сразу переводит текст в числа.

Переполнение

Помни про размер целых чисел: на сайте numpy по умолчанию использует int32 (до ~2·10⁹). Если суммы могут быть больше, читай данные с dtype=np.int64. Обычные числа Python не переполняются никогда, а числа numpy — молча «заворачиваются»:

a = np.array([2_000_000_000, 2_000_000_000], dtype=np.int32)
print(a.sum())                  # может получиться отрицательное число!
print(a.astype(np.int64).sum()) # 4000000000

Итог

• операции и функции numpy (np.sqrt, np.abs, …) применяются к каждому элементу;
• sum, mean, min, max, argmin, argmax, std — агрегаты;
• axis=0 — результат по столбцам, axis=1 — по строкам;
• broadcasting растягивает меньший массив; для строк нужна форма (n, 1);
• для больших сумм используй np.int64.