NumPy vectorization и broadcasting на собеседовании Data Scientist

Закрепи Python для аналитика
200+ задач по pandas, numpy и работе с данными — с разборами
Тренировать Python в Telegram

Карьерник — Duolingo для аналитиков: 10 минут в день тренируй SQL, Python, A/B, статистику, метрики и ещё 3 темы собеса. 1500+ вопросов в Telegram-боте. Бесплатно.

Зачем разбирать на собесе

NumPy — фундамент DS. Без понимания vectorization код медленнее в 100×, без понимания broadcasting — пишутся лишние циклы. На собесе DS: «зачем vectorization», «правила broadcasting», «отличие view от copy».

Главная боль без понимания — DS пишет циклы Python по массиву на 1М элементов, ждёт минуту. Векторизованная версия — миллисекунды.

Что такое vectorization

Идея. Заменить Python-цикл на единственный вызов NumPy-операции, которая выполнится в C-коде с использованием SIMD.

# плохо — Python цикл
result = []
for x in arr:
    result.append(x * 2 + 1)

# хорошо — vectorized
result = arr * 2 + 1

Почему быстрее:

  • Python for интерпретирует каждую итерацию (накладные расходы на bytecode dispatch).
  • NumPy выполняет операцию в C-коде с использованием BLAS / SIMD-инструкций (AVX, AVX-512).
  • Нет boxing/unboxing объектов — данные в контейнере fixed-type.

Когда vectorization работает:

  • Арифметика — +, -, *, /, **.
  • Element-wise — np.exp, np.log, np.sin.
  • Логика — &, |, ~, np.where.
  • Reductions — sum, mean, std, argmax.

Когда не работает:

  • Зависимость текущей итерации от предыдущей (a[i] = f(a[i-1])).
  • Сложные branching, который не выражается через where / маски.
  • Чтение разнотипных JSON / sparse объектов.

Broadcasting: правила

Broadcasting — автоматическое расширение массивов до совместимых форм.

Правила:

  1. Если размерности разные, добавляются единицы слева у меньшего.
  2. Размеры в каждой оси совместимы, если они равны или один из них = 1.
  3. Массив с size=1 в оси «растягивается» виртуально (без копирования).
A = np.array([[1, 2, 3], [4, 5, 6]])    # shape (2, 3)
b = np.array([10, 20, 30])               # shape (3,)
A + b                                     # shape (2, 3) — b расширяется
# [[11, 22, 33],
#  [14, 25, 36]]

c = np.array([[10], [20]])               # shape (2, 1)
A + c                                     # shape (2, 3) — c расширяется
# [[11, 12, 13],
#  [24, 25, 26]]

Не работает:

A = np.zeros((3, 4))    # shape (3, 4)
b = np.zeros(3)          # shape (3,) → (1, 3) — не расширяется до (3, 4)
# A + b → ValueError

Применение в DS:

  • Нормализация — (X - X.mean(axis=0)) / X.std(axis=0) (axis=0 даёт shape (D,), broadcast по строкам).
  • Distance matrix — ((X[:, None, :] - X[None, :, :])**2).sum(-1) ((N,N) матрица расстояний).
  • Маскирование — X[mask] где mask shape совместим.

dtype и память

NumPy хранит данные fixed-type. Тип определяет размер и точность.

np.int8       # 1 byte, [-128, 127]
np.int32      # 4 bytes
np.int64      # 8 bytes (default int)
np.float32    # 4 bytes, ~7 цифр precision
np.float64    # 8 bytes (default float), ~15 цифр
np.bool_      # 1 byte
np.complex64  # 8 bytes

Memory: массив на 1М float32 = 4 МБ. На float64 = 8 МБ.

Performance vs precision:

  • ML inference часто float32 хватает (или INT8 quantized).
  • Финансы — float64 минимум (округление!).
  • Статистика — обычно float64.
arr = np.array([1.0, 2.0], dtype=np.float32)
arr.itemsize    # 4
arr.nbytes      # 8 (2 элемента × 4 byte)

Тонкость: np.float32(1) + np.int64(1) → результат float64 (upcast). Контролируй через explicit astype.

View vs copy

View — два массива ссылаются на одни данные. Изменение одного видно в другом.

Copy — независимая копия данных.

a = np.array([1, 2, 3, 4])
b = a[:2]         # view
b[0] = 99
a                  # array([99, 2, 3, 4]) — изменилось!

c = a[:2].copy()   # copy
c[0] = -1
a                  # без изменения

Что возвращает view:

  • Slicing с шагом — a[::2], a[1:5].
  • Reshape, transpose — a.T, a.reshape(...).
  • Изменение dtype через view.

Что возвращает copy:

  • Fancy indexing — a[[0, 2, 4]].
  • Boolean indexing — a[mask].
  • Арифметика — a + 1.
  • np.copy, astype.

Контроль: arr.flags.owndata — True для copy.

Тонкость: изменения через view иногда выглядят как баг, особенно когда передаёшь массив в функцию и она «случайно» меняет внешний.

Закрепи Python для аналитика
200+ задач по pandas, numpy и работе с данными — с разборами
Тренировать Python в Telegram

Memory layout: C vs Fortran order

Многомерный массив в памяти — линейный буфер. Layout определяет порядок:

  • C order (row-major) — по строкам. a[i, j] лежит соседям с a[i, j+1].
  • Fortran order (column-major) — по столбцам. a[i, j] соседствует с a[i+1, j].
a = np.zeros((3, 4))
a.flags.c_contiguous   # True
a.T.flags.c_contiguous  # False (но F_CONTIGUOUS=True)

Влияние:

  • Прохождение по последней оси в C order — кэш-эффективно.
  • Многие BLAS-вызовы быстрее для C order, некоторые — для Fortran.
  • np.asfortranarray(a) или np.ascontiguousarray(a) для явного контроля.

В DS чаще не критично — pandas и sklearn кладут как удобнее им. Но в numerics (matrix multiplications) layout влияет на скорость в 2-3×.

Когда NumPy медленнее: GIL и Python overhead

GIL отпускается в C-коде. Поэтому threading с numpy-операциями реально параллелится.

Overhead на маленьких операциях. Каждый numpy-вызов имеет ~5-50 мкс overhead. Для 5 элементов — Python list быстрее. Для 5000+ — numpy вне конкуренции.

Список альтернатив:

  • NumPy — стандарт.
  • Numba — JIT-компилятор, ускоряет циклы, особенно полезен для невекторизуемых паттернов.
  • Cython — компиляция в C, ручное ускорение.
  • CuPy — drop-in NumPy на GPU.
  • JAX — NumPy-совместимый API + autograd + XLA-компилятор.
  • PyTorch / TensorFlow — для GPU и автодифференцирования.
  • Polars / DuckDB — для табличных данных вместо pandas.

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

Цикл по строкам DataFrame / array. for i in range(len(arr)) обычно решается векторизацией.

Создавать массив в цикле. arr = np.append(arr, x) каждую итерацию — копирование. Препроцессинг в list, потом np.array(list).

Игнорировать broadcasting. np.tile или np.repeat для расширения — overhead. Broadcasting делает то же без копирования.

.values вместо .to_numpy(). В новых pandas .values deprecated. Плюс to_numpy() контролирует dtype.

Большой массив дефолтного типа. 1М objects array (string, dict) — Python-overhead огромный. Если нужны числа — fixed dtype.

Изменять view, не подозревая. Передача arr[:5] в функцию, функция меняет → внешний массив тоже меняется.

np.float (deprecated). Использовать np.float64 или Python float.

Сравнение float через ==. Использовать np.isclose(a, b, rtol=1e-5).

Связанные темы

FAQ

NumPy быстрее pandas?

NumPy на ndarray быстрее pandas DataFrame на простых операциях (там нет overhead на индексы / labels). Pandas удобнее для табличных задач, NumPy — для numerics.

Когда использовать np.einsum?

Для сложных tensor-операций, которые не выражаются стандартными ops. Например, np.einsum('ij,jk->ik', A, B) = A @ B. Гибче, но иногда медленнее, чем явные операции.

Какой dtype для категориальных?

np.int8 или np.int16 если категорий меньше 32k. Pandas Categorical или int8 дешевле, чем object (string).

np.where vs maska?

np.where(cond, x, y) — выбор между двумя массивами по условию. Маска arr[mask] — выборка. Разные задачи.

Что такое structured array?

Массив с смешанными типами (как DataFrame). Реально используется редко — чаще берут pandas / pyarrow. NumPy structured подходит для interop с C.

Это официальная информация?

Нет. Статья основана на документации NumPy 1.26+ и материалах из «Effective Python».


Тренируйте Data Science — откройте тренажёр с 1500+ вопросами для собесов.