Собесов

Сценарий: numpy broadcasting — посчитать матрицу попарных расстояний без цикла

PythonNumpy и scipyСредняяMiddle

Условие

Дана матрица 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))   # та же матрица, чуть быстрее

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

  1. (n, n, d) для n=10000, d=100 = ~80 ГБ — наивный broadcasting убьёт память. Использовать тождество или cdist.
  2. Из-за арифметики float64 диагональ может оказаться -1e-15 — после sqrt будет nan. Защита: np.maximum(D2, 0).
  3. X[:, None]X[None, :] — путаница порядков осей даёт неверный broadcasting.
  4. cdist(X, X) симметрична, считается ровно в 2 раза дольше нужного. pdist + squareform экономит.
  5. Для огромных n — chunk-based вычисление или приближение (faiss/annoy).

Эталонный ответ

Через тождество ||x-y||² = ||x||² + ||y||² − 2X·Xᵀ. Готовое — scipy.spatial.distance.cdist.

Хочешь увидеть разбор?

Зарегистрируйся бесплатно — откроется развёрнутое решение этой задачи и ещё 4 на выбор.

Зарегистрироваться и увидеть разбор
Уже есть аккаунт? Войти