NumPy за 7 минут: массивы, broadcasting, векторизация

pandas построен поверх NumPy, и половина вопросов про «почему pandas так себя ведёт» на самом деле про NumPy. Главная идея одна: массив хранит числа одного типа подряд в памяти, а операции над ним выполняются циклами на C.

ndarray: форма, тип, оси

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6]])
a.shape   # (2, 3): 2 строки, 3 столбца
a.ndim    # 2
a.dtype   # int64
a.size    # 6

У массива один dtype на все элементы. Если записать в целочисленный массив 2.7, получится 2: значение молча приведётся к типу массива. Строки и числа вперемешку дадут массив строк, а не «умный» список.

Ось (axis) это направление, вдоль которого идёт агрегат. axis=0 схлопывает строки, результат по одному числу на столбец; axis=1 схлопывает столбцы.

a.sum(axis=0)   # [5 7 9]
a.sum(axis=1)   # [ 6 15]
a.mean()        # 3.5, по всем элементам

Векторизация

Цикл по массиву в Python медленный: на каждом шаге интерпретатор достаёт элемент, упаковывает его в объект и проверяет типы. Векторная операция делает то же самое внутри скомпилированного кода.

x = np.arange(1_000_000)
y = x * 2 + 1                 # быстро, без цикла Python
y = np.where(x % 2 == 0, x, -x)   # if-else поэлементно

Инструменты векторизации: арифметика массивов, универсальные функции (np.exp, np.log, np.abs), сравнения и булевы маски, np.where, np.select, агрегаты с axis, np.cumsum, np.diff, np.sort, np.argsort.

np.vectorize к ним не относится: это удобная обёртка, которая всё равно вызывает вашу Python-функцию на каждом элементе.

Broadcasting

Broadcasting позволяет делать операции над массивами разной формы без копирования. Правила:

1. Формы выравниваются по правому краю, недостающие размерности слева считаются равными 1.
2. Размеры по каждой оси должны совпадать или один из них должен быть 1.
3. Ось размера 1 «растягивается» до размера другой.

row = np.array([1, 2, 3])        # (3,)
col = np.array([[10], [20]])     # (2, 1)
row + col                        # (2, 3): [[11 12 13], [21 22 23]]

Типичное применение: вычесть из матрицы средние по столбцам X - X.mean(axis=0) или посчитать все попарные разности p[:, None, :] - p[None, :, :]. Если формы несовместимы, будет ValueError: operands could not be broadcast together.

Индексация: view и copy

Это самая частая ловушка.

• Базовый срез (a[1:3], a[:, 0], a[::2]) возвращает представление (view): новый объект смотрит в ту же память. Изменили срез, изменился исходный массив.
• Fancy indexing (список или массив индексов a[[0, 2]]) и булева маска (a[a > 0]) всегда создают копию.

a = np.arange(6)
b = a[::2]        # view
b += 10           # a тоже изменился: [10 1 12 3 14 5]
c = a[[1, 3]]     # copy
c += 100          # a не меняется

Проверить, делят ли массивы память, можно через np.shares_memory(a, b). Если собираетесь менять данные и не хотите задеть исходник, делайте .copy() явно. Ещё одна неочевидная деталь: маленький срез держит ссылку на весь большой массив, и память не освобождается.

reshape и .T тоже стараются вернуть view, но если новую форму нельзя описать шагами по текущей памяти (например, сплющить транспонированный массив), reshape молча скопирует.

Типы и переполнение

Целые типы фиксированной ширины переполняются без ошибок, по модулю:

x = np.array([200], dtype=np.uint8)
x + x   # array([144], dtype=uint8): 400 - 256

С пропусками в целых NumPy не работает: NaN бывает только во float. Сравнение с NaN всегда ложно, поэтому ищите пропуски через np.isnan, а агрегаты без них считайте через np.nanmean, np.nansum.

Полезные приёмы

• np.argsort возвращает индексы сортировки, np.argmax индекс максимума, np.argsort(-a)[:k] даёт топ-k.
• Скользящая сумма окна: c = np.cumsum(a) с нулём в начале, окно [i, i+k) это c[i+k] - c[i].
• np.percentile(a, [5, 95]) и np.clip(a, lo, hi) обрезают выбросы.
• np.random.default_rng(42) современный генератор случайных чисел с воспроизводимым seed.

Частые ошибки

• Ожидать, что срез NumPy это копия, как у списков Python.
• Путать axis=0 и axis=1.
• Писать цикл for по массиву или использовать np.vectorize ради скорости.
• Сравнивать с np.nan через ==.
• Копить суммы в маленьком целом типе и ловить переполнение.

Что спросят на собесе

• Почему NumPy быстрее списков и что такое векторизация?
• Правила broadcasting: какая форма будет у (3, 1) + (4,)?
• Какие операции возвращают view, а какие copy? Как проверить?
• Что значит axis=0 в агрегате?
• Как без цикла посчитать скользящее среднее или попарные расстояния?