记矩阵 $A$ 大小为 $M\times D$，矩阵 $B$ 大小为 $N\times D$，$A=[{\vec{a}}_1,{\vec{a}}_2\cdots {\vec{a}}_m]$，$B=[{\vec{b}}_1,{\vec{b}}_2\cdots {\vec{b}}_n]$
现在要计算 $A$ 中每个向量 ${\vec{a}}_i$ 与 $B$ 中每个向量 ${\vec{b}}_j$ 之间的欧式距离，最朴素的办法就是 $O(n^2)$ 的硬算，但这样会非常的慢，如果可以使用矩阵运算来代替循环，则可以大大提升速度。

推导过程

首先计算 ${\vec{a}}_i$ 和 ${\vec{b}}_j$ 之间的距离 $D(i,j)$
$$\begin{array}{rl}D(i,j)& =\sqrt{(a_{i,1}-b_{j,1}{)}^2+\cdots +(a_{i,D}-b_{j,D}{)}^2}\\ & =\sqrt{(a_{i,1}^2+\cdots +a_{i,D}^2)+(b_{j,1}^2+\cdots +b_{j,D}^2)-2\times (a_{i,1}b{+}_{j,1}+\cdots +a_{i,D}+b_{j,D})}\\ & =\sqrt{||{\vec{a}}_i|{|}^2+||{\vec{b}}_j|{|}^2-2\times {\vec{a}}_i{\vec{b}}_j^T}\end{array}$$

由此再推广至整个矩阵
$$\text{dist}=\sqrt{{\left(\begin{array}{ccc}||{\vec{a}}_1|{|}^2& \cdots & ||{\vec{a}}_1|{|}^2\\ ⋮& \ddots & ⋮\\ ||{\vec{a}}_M|{|}^2& \cdots & ||{\vec{a}}_M|{|}^2\end{array}\right)}_{M\times N}+{\left(\begin{array}{ccc}||{\vec{b}}_1|{|}^2& \cdots & ||{\vec{b}}_N|{|}^2\\ ⋮& \ddots & ⋮\\ ||{\vec{b}}_1|{|}^2& \cdots & ||{\vec{b}}_N|{|}^2\end{array}\right)}_{M\times N}-2\times A{B}^T}$$

代码

def euclidean_dist(x, y):
    """
    PARAMETER:
    x: pytorch Variable, with shape [m, d]
    y: pytorch Variable, with shape [n, d]
    RETURN:
    dist: pytorch Variable, with shape [m, n]
    """

    m, n = x.size(0), y.size(0)
    # x 对每个数进行平方后，在 axis=1 方向加和，此时 x shape 为 (m, 1)，经过 expand() 扩展为 (m, n)
    xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
    yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
    dist = xx + yy
    # help(torch.addmm())
    # help(dist.addmm())
    torch.addmm(dist, x, y.t(), beta=1, alpha=-2, out=dist)
    return dist.sqrt()