如何调整代码实现特征同维、样本数不同数组的欧氏距离计算

问题根源 #

原有代码的问题集中在两个点：

• 计算S、R时硬编码了repeat的次数（3、2），这两个值是测试用小数据集的样本量，换成任意尺寸的X、Z就会出现维度不匹配报错
• 全局变量耦合：l2distance函数直接使用外部定义的S、R、G，没有和输入的X、Z动态绑定，复用性差

修正后的完整代码 #

import numpy as np

def l2distance(X,Z=None):
    if Z is None:
        Z = X
    n, d1 = X.shape
    m, d2 = Z.shape
    assert d1 == d2, "X和Z的特征维度必须一致"
    
    # 计算G：X和Z的内积，形状(n, m)
    G = np.inner(X, Z)
    
    # 计算S：X每个样本自身内积，拓展为(n, m)
    S = np.diag(np.inner(X, X))[:, None].repeat(m, axis=1)
    
    # 计算R：Z每个样本自身内积，拓展为(n, m)
    R = np.diag(np.inner(Z, Z))[None, :].repeat(n, axis=0)
    
    return S + R - 2 * G

# 小样本测试验证
X = np.array([[1,2],[3,4]])
Z = np.array([[1,4],[2,5],[3,6]])
print("小样本测试结果：\n", l2distance(X,Z))

# 大尺寸输入测试
X_big = np.random.rand(700,100)
Z_big = np.random.rand(800,100) 
Dsqr_big = l2distance(X_big, Z_big)
print("大样本输出形状：", Dsqr_big.shape) # 输出应为(700, 800)符合预期

关键调整说明 #

• 移除了硬编码的repeat次数：S重复次数替换为Z的样本量m，R重复次数替换为X的样本量n，适配任意输入尺寸
• 把S、R、G的计算全部整合到l2distance函数内部，消除全局变量依赖，每次调用都会根据输入的X、Z动态计算对应值
• 如需进一步优化性能，可以把repeat操作替换为numpy广播计算：S = np.sum(X**2, axis=1, keepdims=True)、R = np.sum(Z**2, axis=1, keepdims=True).T，效果完全一致，大矩阵场景下计算效率更高

内容的提问来源于stack exchange，提问作者Slavisha84