Библиотека NumPy для Python: основы работы с массивами

Библиотека NumPy для Python дает быстрый многомерный массив ndarray и сотни функций для вычислений над ним без циклов. Ниже по шагам разобраны установка, создание массивов, индексация и маски, математические операции, параметр axis, broadcasting, линейная алгебра и генерация случайных чисел.

Установка и импорт

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

pip install numpy

import numpy as np
print(np.__version__)  # например 2.1.3

NumPy лежит в основе pandas, Matplotlib, scikit-learn и SciPy, поэтому при установке этих библиотек он подтягивается автоматически. Вторая мажорная версия вышла в 2024 году, в ней удалены старые псевдонимы типов вроде np.float и np.int: вместо них пишут встроенные float, int или np.float64.

Создание массивов

Главный объект библиотеки - массив ndarray. В отличие от списка, все его элементы имеют один тип и лежат в памяти подряд, поэтому операции над ними выполняются в скомпилированном коде на C.

import numpy as np

a = np.array([1, 2, 3, 4, 5])
zeros = np.zeros((2, 3))         # матрица 2x3 из нулей (float)
ones = np.ones(3, dtype=int)     # [1 1 1]
r = np.arange(0, 10, 2)          # [0 2 4 6 8]
l = np.linspace(0, 1, 5)         # [0.   0.25 0.5  0.75 1.  ]
m = np.arange(1, 7).reshape(2, 3)
print(m)
# [[1 2 3]
#  [4 5 6]]

arange() работает как range(), но возвращает массив и принимает дробный шаг. linspace() делит отрезок на заданное число точек, включая оба конца. Метод reshape() меняет форму без копирования данных, произведение размеров должно совпадать с числом элементов.

У каждого массива есть атрибуты, которые нужно знать с первого дня:

Индексы, срезы и логические маски

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

x = np.array([10, 20, 30, 40, 50])
print(x[0], x[-1])    # 10 50
print(x[1:4])         # [20 30 40]

m = np.arange(1, 7).reshape(2, 3)
print(m[1, 2])        # 6
print(m[:, 0])        # первый столбец: [1 4]

print(x[x > 25])                 # [30 40 50]
print(np.where(x > 25, x, 0))    # [ 0  0 30 40 50]

Выражение x > 25 дает массив из True и False, его называют маской. Подстановка маски в квадратные скобки оставляет только подходящие элементы, это быстрая замена цикла с условием. Условия объединяют операторами & и |, каждое в скобках: x[(x > 10) & (x < 50)].

Важное отличие от списков: срез массива - это представление (view), а не копия. Изменение среза меняет исходный массив.

x = np.array([1, 2, 3, 4, 5])
s = x[1:4]
s[0] = 99
print(x)  # [ 1 99  3  4  5]
# независимая копия: s = x[1:4].copy()

Математика без циклов

Арифметические операторы применяются к каждому элементу сразу, такой стиль называют векторизацией. Код получается короче, а работает в десятки раз быстрее цикла for по списку.

x = np.array([1, 2, 3, 4, 5])
print(x * 2)        # [ 2  4  6  8 10]
print(x ** 2)       # [ 1  4  9 16 25]
print(x + x)        # [ 2  4  6  8 10]
print(np.sqrt(x))   # корень из каждого элемента
print(x.sum(), x.mean(), x.max(), x.std())

У многомерных массивов агрегатные функции принимают параметр axis. Значение axis=0 сворачивает строки и считает результат по каждому столбцу, axis=1 - по каждой строке.

m = np.array([[1, 2, 3],
              [4, 5, 6]])
print(m.sum())          # 21
print(m.sum(axis=0))    # [5 7 9]
print(m.sum(axis=1))    # [ 6 15]

Broadcasting: операции с массивами разной формы

Если формы массивов не совпадают, NumPy пытается растянуть меньший массив по недостающей оси. Правило такое: размеры сравниваются справа налево и считаются совместимыми, если равны или один из них равен 1.

m = np.array([[1, 2, 3],
              [4, 5, 6]])
row = np.array([10, 20, 30])
print(m + row)
# [[11 22 33]
#  [14 25 36]]

Строка из трех элементов прибавилась к каждой строке матрицы, при этом реальная копия не создавалась. При несовместимых формах, например (2, 3) и (2,), будет ошибка ValueError: operands could not be broadcast together.

Линейная алгебра и случайные числа

Матричное умножение выполняет оператор @, а звездочка * умножает поэлементно. Модуль np.linalg находит определитель, обратную матрицу и решает системы линейных уравнений.

A = np.array([[2, 1],
              [1, 3]])
b = np.array([3, 5])
print(A @ A)                  # [[ 5  5] [ 5 10]]
print(A.T)                    # транспонирование
print(np.linalg.solve(A, b))  # [0.8 1.4]

Система 2x + y = 3, x + 3y = 5 решена одной строкой. Для случайных чисел современный способ - генератор default_rng(), функции старого интерфейса np.random.rand() работают, но в новом коде их не рекомендуют.

rng = np.random.default_rng(42)  # зерно для повторяемости
print(rng.integers(1, 7, size=5))  # пять бросков кубика
print(rng.normal(0, 1, size=3))    # нормальное распределение

NumPy или обычный список

Список Python хранит ссылки на объекты любых типов и легко растет через append(). Массив NumPy имеет фиксированный размер и один тип, зато расходует меньше памяти и быстро считает. Отсюда практическое правило:

• числовые данные, над которыми нужно считать, - массив NumPy;
• разнородные объекты, частое добавление и удаление элементов - список;
• таблицы с подписанными столбцами и пропусками - pandas, который построен поверх NumPy.

Добавлять элементы в массив в цикле через np.append() медленно: функция каждый раз создает новый массив. Лучше собрать данные в список и один раз превратить его в массив. Полный справочник функций есть в официальной документации NumPy.

Частые вопросы

Ответы на вопросы, с которыми сталкиваются при первом знакомстве с NumPy.

Как превратить массив NumPy обратно в список?

Вызовите метод tolist(). Он вернет обычный список, для многомерного массива вложенный, а элементы станут встроенными типами int и float.

Почему 0.1 + 0.2 в массиве не равно 0.3?

NumPy использует те же числа с плавающей точкой, что и Python. Для сравнения дробных значений применяйте np.isclose(a, b) или np.allclose() для целых массивов.

Что значит ошибка ModuleNotFoundError: No module named numpy?

Пакет не установлен в тот интерпретатор, которым запускается программа. Выполните python -m pip install numpy тем же Python, что запускает скрипт, или активируйте нужное виртуальное окружение.

Чем shape (3,) отличается от (3, 1)?

Форма (3,) означает одномерный массив из трех чисел. Форма (3, 1) - это двумерная матрица из трех строк и одного столбца. Перевести одно в другое можно через reshape(-1, 1).