如何无需for循环，用Numpy实现二维数组逐行批量相减？

问题背景 #

在Numpy中，当二维数组A（形状5000×3072）与行向量v（形状1×3072）相减时，因列数匹配，v会自动广播为与A同形状的数组完成运算。但如果要让二维数组V的每一行分别与整个A做运算，直接操作无法实现。现有一个带for循环的半向量化距离计算函数，需要改造为完全无循环的实现。

现有半向量化代码 #

def compute_distances_one_loop(V):
    num_test = V.shape[0]
    num_train = A.shape[0]
    dists = np.zeros((num_test, num_train))
    for i in range(num_test):   
        initial = np.sqrt(np.square(A - V[i,:]))
        dists[i,:] = initial.sum(axis=1)
    return dists

示例场景 #

给定以下数组：

matrix_e = np.ones((3,3)) * 6
matrix_f = np.array([[1,2,3],[4,5,6]])

如何无需for循环得到形状为(6,3)的matrix_g：

matrix_g = np.array([[5,4,3],[5,4,3],[5,4,3],[2,1,0],[2,1,0],[2,1,0]])

解决方案 #

1. 完全向量化的距离计算 #

原函数的核心是计算V的每一行与A的每一行对应元素的绝对值之和（注：np.sqrt(np.square(x))等价于np.abs(x)）。利用Numpy的广播机制，通过扩展维度实现无循环运算：

def compute_distances_no_loop(V):
    # 扩展A的维度为(1, num_train, num_features)
    A_expanded = A[np.newaxis, :, :]
    # 扩展V的维度为(num_test, 1, num_features)
    V_expanded = V[:, np.newaxis, :]
    # 广播相减后计算绝对值，再沿特征维度求和
    dists = np.abs(A_expanded - V_expanded).sum(axis=2)
    return dists

扩展维度后，A_expanded与V_expanded相减会自动广播为(num_test, num_train, num_features)的数组，最后沿第三维度（特征维度）求和，直接得到num_test×num_train的距离矩阵。

2. 示例场景的无循环实现 #

要得到目标matrix_g，本质是让matrix_f的每一行分别与matrix_e的所有行相减，再将结果展平为6×3的数组。两种实现方式：

方式一：广播+维度重塑

# 扩展matrix_f为(2,1,3)，matrix_e为(1,3,3)，广播相减得到(2,3,3)
temp = matrix_e - matrix_f[:, np.newaxis, :]
# 重塑为(6,3)
matrix_g = temp.reshape(-1, 3)

方式二：重复数组+直接运算

# 将matrix_f的每一行重复3次，得到(6,3)
f_repeated = np.repeat(matrix_f, repeats=3, axis=0)
# 将matrix_e重复2次，得到(6,3)
e_repeated = np.tile(matrix_e, (2,1))
matrix_g = e_repeated - f_repeated

内容的提问来源于stack exchange，提问作者Ülke Eren Aktaş