Условие
Дана матрица X размера (n, d) — n точек в d-мерном пространстве. Посчитать матрицу попарных евклидовых расстояний D размера (n, n) без Python-цикла.
Решение
Идея броадкастинга
X[:, None, :] имеет форму (n, 1, d), X[None, :, :] — (1, n, d). При вычитании broadcasting расширит их до (n, n, d). Дальше — sum по последней оси.
Реализация
import numpy as np
def pairwise_dist(X):
# (n, 1, d) - (1, n, d) -> (n, n, d)
diff = X[:, None, :] - X[None, :, :]
sq = (diff ** 2).sum(axis=-1)
return np.sqrt(sq)Более эффективно через тождество
||x - y||² = ||x||² + ||y||² − 2x·y — избегает создания (n, n, d)-тензора.
def pairwise_dist_fast(X):
sq = (X ** 2).sum(axis=1)
G = X @ X.T
D2 = sq[:, None] + sq[None, :] - 2 * G
D2 = np.maximum(D2, 0) # численная стабильность
return np.sqrt(D2)При больших n второй вариант экономит в d раз память (нет (n,n,d)-тензора).
Готовое решение
from scipy.spatial.distance import cdist, squareform, pdist
D = cdist(X, X) # (n, n)
D = squareform(pdist(X)) # та же матрица, чуть быстрееПодводные камни
(n, n, d)для n=10000, d=100 = ~80 ГБ — наивный broadcasting убьёт память. Использовать тождество илиcdist.- Из-за арифметики float64 диагональ может оказаться
-1e-15— послеsqrtбудетnan. Защита:np.maximum(D2, 0). X[:, None]≠X[None, :]— путаница порядков осей даёт неверный broadcasting.cdist(X, X)симметрична, считается ровно в 2 раза дольше нужного.pdist+squareformэкономит.- Для огромных n — chunk-based вычисление или приближение (faiss/annoy).
Эталонный ответ
Через тождество ||x-y||² = ||x||² + ||y||² − 2X·Xᵀ. Готовое — scipy.spatial.distance.cdist.