Нашли ошибку?
продвинутыйNumPy 2.5.3

Векторизация вместо циклов

Как переписать цикл Python на операции над целыми массивами и почему np.vectorize не делает код быстрее.

Нужно знать:where, nonzero, selectТранслирование (broadcasting)Универсальные функции (ufunc)

Содержание статьи

Коротко

  • Векторизация — запись вычисления через операции над целыми массивами. Цикл выполняет NumPy на C, а не интерпретатор Python — обычно в 10–100 раз быстрее.
  • Типовые замены: if/else → np.where, цепочка if/elif → np.select, накопление суммы → cumsum, двойной цикл «каждый с каждым» → транслирование, подсчёт по группам → bincount.
  • np.vectorize только оборачивает функцию Python, чтобы её можно было вызвать от массива. Внутри тот же цикл по элементам, ускорения нет.
  • Векторизованный код иногда создаёт большие промежуточные массивы — следите за памятью.

Синтаксис

# цикл                                   # векторно
[f(x) if c else g(x) for x in a]         np.where(cond, f(a), g(a))
if / elif / else                         np.select([c1, c2], [v1, v2], default)
s += x в цикле                           np.cumsum(a), a.sum()
for i: for j: d[i, j] = …               a[:, None] - a[None, :]
np.vectorize(f)(a)                       # удобно, но не быстрее цикла

Примеры

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

rng = np.random.default_rng(0)
amounts = rng.uniform(100, 5000, size=100_000)   # суммы покупок

def bonus_vector(values):                        # бонус 5 % с покупок дороже 1000
    return np.where(values > 1000, values * 0.05, 0.0)
Цикл и векторная версия
def bonus_loop(values):
    result = []
    for v in values:
        result.append(v * 0.05 if v > 1000 else 0.0)
    return np.array(result)

print(np.allclose(bonus_loop(amounts), bonus_vector(amounts)))   # результат одинаковый, а время — нет
t_loop = min(timeit.repeat(lambda: bonus_loop(amounts), number=1, repeat=5))
t_vector = min(timeit.repeat(lambda: bonus_vector(amounts), number=1, repeat=5))
print(f"цикл: {t_loop * 1000:.1f} мс, векторно: {t_vector * 1000:.2f} мс")
print(f"ускорение: ×{t_loop / t_vector:.0f}")
Выводзамер на Apple M4 · macOS · Python 3.14, 29.09.2026
True
цикл: 6.4 мс, векторно: 0.06 мс
ускорение: ×112

Время снято один раз и показано вместе с машиной и датой замера. На другом компьютере числа будут другими, но порядок разницы сохранится.

Накопление — cumsum
# «Сколько потрачено к каждой покупке» — не цикл с суммой, а cumsum
day = np.array([300.0, 150.0, 700.0, 50.0])
np.cumsum(day)
Вывод
array([ 300.,  450., 1150., 1200.])
if / elif — np.select
# цикл с if/elif превращается в np.select
tier = np.select([amounts < 500, amounts < 2000], ["малая", "средняя"], default="крупная")
tier[:5]
Вывод
array(['крупная', 'средняя', 'малая', 'малая', 'крупная'], dtype='<U7')
Двойной цикл — транслирование
# двойной цикл «каждый с каждым» — транслирование
points = np.array([[0.0, 0.0], [3.0, 4.0], [6.0, 8.0]])
diff = points[:, np.newaxis, :] - points[np.newaxis, :, :]
np.sqrt((diff ** 2).sum(axis=-1))                                 # матрица расстояний
Вывод
array([[ 0.,  5., 10.],
       [ 5.,  0.,  5.],
       [10.,  5.,  0.]])

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

np.vectorize не ускоряет

Название обманчиво: это удобная обёртка для вызова обычной функции на массиве, а не компиляция. Если функцию нельзя переписать через операции NumPy, np.vectorize хотя бы делает код короче — но не быстрее.

np.vectorize против векторной версии
def bonus_one(v):
    return v * 0.05 if v > 1000 else 0.0

bonus_vectorized = np.vectorize(bonus_one)
print(np.allclose(bonus_vectorized(amounts), bonus_vector(amounts)))
t_vectorize = min(timeit.repeat(lambda: bonus_vectorized(amounts), number=1, repeat=5))
t_vector = min(timeit.repeat(lambda: bonus_vector(amounts), number=1, repeat=5))
print(f"np.vectorize: {t_vectorize * 1000:.1f} мс, np.where: {t_vector * 1000:.2f} мс")
print(f"np.where быстрее: ×{t_vectorize / t_vector:.0f}")   # np.vectorize — тот же цикл внутри
Выводзамер на Apple M4 · macOS · Python 3.14, 29.09.2026
True
np.vectorize: 5.8 мс, np.where: 0.06 мс
np.where быстрее: ×101