NumPy: дана матрица точек X размера (n, d) и матрица центров C размера (k, d). Посчитайте без циклов матрицу попарных евклидовых расстояний и ближайший центр для каждой точки.
Короткий ответ
- Броадкастинг: X[:, None, :] минус C[None, :, :]
- Квадрат, сумма по последней оси, корень
- Ближайший центр — argmin по оси центров
- Альтернатива через (a-b)^2 = a^2 - 2ab + b^2 экономит память
- Векторизация на порядки быстрее питоновских циклов
- Это ядро шага присвоения в k-means
Задача решается броадкастингом с разницей по новой оси и argmin, а для больших данных — разложением квадрата разности без материализации тензора (n, k, d).
Как сказать вслух
пример ответаЯ использую броадкастинг: добавляю точкам и центрам по новой оси, чтобы при вычитании получился тензор эн на ка на дэ с попарными разностями. Дальше возвожу в квадрат, суммирую по последней оси и беру корень — получается матрица расстояний, а argmin по оси центров даёт ближайший центр. Если данных много, вместо огромного промежуточного тензора раскладываю квадрат разности по формуле — память экономится в разы.
Подробный ответ
Основной ответ
Идиоматичное решение — броадкастинг: diff = X[:, None, :] - C[None, :, :] даёт тензор (n, k, d), затем dist = np.sqrt((diff ** 2).sum(axis=-1)) — матрицу (n, k), а labels = dist.argmin(axis=1) — индекс ближайшего центра. Это шаг присвоения из k-means. Нюанс для senior-уровня — память: тензор (n, k, d) при n=10^6 может не влезть, и тогда используют разложение |a-b|^2 = |a|^2 - 2ab + |b|^2: квадраты норм считаются отдельно, перекрёстный член — одним матричным умножением X @ C.T, и промежуточный объект имеет размер лишь (n, k). Для поиска ближайшего центра корень извлекать не обязательно — argmin монотонен по квадрату расстояния. Векторизация переносит вычисления в оптимизированный C-код и ускоряет решение на один-два порядка против вложенных циклов.
Ключевые моменты
- Броадкастинг через новые оси. X[:, None] и C[None, :] выравнивают формы (n,1,d) и (1,k,d) для попарного вычитания.
- Экономия памяти. Разложение квадрата разности заменяет тензор (n,k,d) матричным умножением и формой (n,k).
- Корень не нужен для argmin. Монотонные преобразования не меняют позицию минимума — мелкая, но показательная оптимизация.
- Связь с k-means. Это шаг assignment кластеризации; в проде та же логика есть в scipy.spatial.distance.cdist.
Практический контекст
Задача проверяет мышление массивами — ключевой навык для ML-инженера: тот, кто пишет вложенные циклы по numpy-массивам, будет тормозить любой пайплайн. Интервьюеры любят усложнение «а если миллион точек?» — ожидая разговор о памяти и разложении формулы. Та же техника используется в kNN, k-means и расчёте матриц близости эмбеддингов.
Пример кода
import numpy as np
def nearest_centers(X: np.ndarray, C: np.ndarray):
# Вариант 1: броадкастинг, тензор (n, k, d)
diff = X[:, None, :] - C[None, :, :]
dist = np.sqrt((diff ** 2).sum(axis=-1)) # (n, k)
# Вариант 2: экономия памяти, промежуточно только (n, k)
sq = (X ** 2).sum(1)[:, None] - 2 * X @ C.T + (C ** 2).sum(1)[None, :]
dist2 = np.sqrt(np.maximum(sq, 0))
labels = dist2.argmin(axis=1) # ближайший центр
return dist2, labels
X = np.random.rand(1000, 16)
C = np.random.rand(8, 16)
dist, labels = nearest_centers(X, C)Частые ошибки
- Пишут двойной цикл по точкам и центрам вместо броадкастинга
- Не думают о памяти: тензор (n, k, d) на больших n приводит к MemoryError
- Путают оси в sum и argmin и получают расстояния неправильной формы