pytorch和Numpy 的广播机制高效计算矩阵之间两两距离

示例：

import torch

def pairwise_distance_broadcast(X, Y=None):
    """
    使用广播机制计算两两欧几里得距离。
    参数:
        X: 第一个张量，形状为 (m, d)，表示 m 个 d 维向量。
        Y: 第二个张量（可选），形状为 (n, d)。如果为 None，则计算 X 内部的两两距离。
    返回:
        距离矩阵，形状为 (m, n)。
    """
    if Y is None:
        Y = X

    # 扩展维度以支持广播
    X_expanded = X.unsqueeze(1)  # 形状为 (m, 1, d)
    Y_expanded = Y.unsqueeze(0)  # 形状为 (1, n, d)

    # 计算两两距离
    distance_matrix = torch.sqrt(torch.sum((X_expanded - Y_expanded) ** 2, dim=2))  # 形状为 (m, n)
    return distance_matrix

# 示例
X = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
Y = torch.tensor([[7.0, 8.0], [9.0, 10.0]], dtype=torch.float32)

# 计算 X 和 Y 之间的两两距离
distances = pairwise_distance_broadcast(X, Y)
print("X 和 Y 之间的两两距离：")
print(distances)

# 计算 X 内部的两两距离
internal_distances = pairwise_distance_broadcast(X)
print("X 内部的两两距离：")
print(internal_distances)

 

 

示例：

'''求三维空间中两组点之间的两两距离'''
import numpy as np 

# 第一组点，假设有四个点，位置如下
a = np.array([[ 0, 0, 0],
                [10,10,10],
                [20,20,20],
                [30,30,30]])
# 第二组点，假设有两个点，位置如下
b = np.array([[ 0, 0, 0],
                [10,10,10]])

# 计算距离
c = a[:, None, :] - b   
# 用None扩展维度，使其shape变为(4, 1, 3)，b的shape是(2, 3)，
# 然后两个数组相减，此时会用到广播机制。 c的shape:(4, 2, 3)
dist = np.sqrt(np.sum(c * c, axis=-1))
print(dist)
# 输出如下：
# [[ 0.         17.32050808]
#  [17.32050808  0.        ]
#  [34.64101615 17.32050808]
#  [51.96152423 34.64101615]]

广播机制的规则可参考：nupmy broadcast (广播) 

 

======================================================================

以下是旧内容：

利用numpy可以很方便的计算两个二维数组之间的距离。二维数组之间的距离定义为：X的维度为(m, c),Y的维度为(n,c),Z为X到Y的距离数组,维度为(m,n)。且Z[0,0]是X[0]到Y[0]的距离。Z(a,b)为X[a]到Y[b]的距离。

例如： 计算 m*2 的矩阵与 n * 2 的矩阵中，m*2 的每一行到 n*2 的两两之间欧氏距离。

'''
L2 = sqrt((x1-x2)^2 + (y1-y2)^2 + (z1-z2)^2)
'''

import numpy as np

# ============= 方法一：不用循环 ====================
def L2_dist_1(cloud1, cloud2):
    m, n = len(cloud1), len(cloud2)
    cloud1 = np.repeat(cloud1, n, axis=0)
    cloud1 = np.reshape(cloud1, (m, n, -1))
    dist = np.sqrt(np.sum((cloud1 - cloud2)**2, axis=2))
    return dist


# ============= 方法二：用一重循环 ==================
def L2_dist_2(cloud1, cloud2):
    m, n = len(cloud1), len(cloud2)
    dist = np.zeros((m, n), dtype=np.float)
    for i in range(m):
        dist[i, :] = np.sqrt(np.sum((cloud1[i, :] - cloud2)**2, axis=1)) 
    return dist


# ============= 方法二：用两重循环 ==================
def L2_dist_3(cloud1, cloud2):
    m, n = len(cloud1), len(cloud2)
    dist = np.zeros((m, n), dtype=np.float)
    for i in range(m):
        for j in range(n):
            dist[i, j] = np.sqrt(np.sum((cloud1[i, :] - cloud2[j, :])**2, axis=0)) 
    return dist


if __name__ == '__main__':
    a = np.array([[ 0, 0, 0],
                  [10,10,10],
                  [20,20,20],
                  [30,30,30]])

    b = np.array([[ 0, 0, 0],
                  [10,10,10]])

    print('不用循环：\n', L2_dist_1(a, b))
    print('用一重循环：\n', L2_dist_2(a, b))
    print('用两重循环：\n', L2_dist_3(a, b))

程序运行结果：

 

参考：Numpy 的广播机制高效计算矩阵之间两两距离